aws-bootstrap-g4dn 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

aws_bootstrap/cli.py

@@ -10,8 +10,13 @@ import click
 
 from .config import LaunchConfig
 from .ec2 import (
+    EBS_MOUNT_POINT,
     CLIError,
+    attach_ebs_volume,
+    create_ebs_volume,
+    delete_ebs_volume,
     ensure_security_group,
+    find_ebs_volumes_for_instance,
     find_tagged_instances,
     get_latest_ami,
     get_spot_price,
@@ -19,13 +24,18 @@ from .ec2 import (
     list_amis,
     list_instance_types,
     terminate_tagged_instances,
+    validate_ebs_volume,
     wait_instance_ready,
 )
+from .output import OutputFormat, emit, is_text
 from .ssh import (
     add_ssh_host,
+    cleanup_stale_ssh_hosts,
+    find_stale_ssh_hosts,
     get_ssh_host_details,
     import_key_pair,
     list_ssh_hosts,
+    mount_ebs_volume,
     private_key_path,
     query_gpu_info,
     remove_ssh_host,
@@ -39,22 +49,32 @@ SETUP_SCRIPT = Path(__file__).parent / "resources" / "remote_setup.sh"
 
 
 def step(number: int, total: int, msg: str) -> None:
+    if not is_text():
+        return
     click.secho(f"\n[{number}/{total}] {msg}", bold=True, fg="cyan")
 
 
 def info(msg: str) -> None:
+    if not is_text():
+        return
     click.echo(f"  {msg}")
 
 
 def val(label: str, value: str) -> None:
+    if not is_text():
+        return
     click.echo(f"  {label}: " + click.style(str(value), fg="bright_white"))
 
 
 def success(msg: str) -> None:
+    if not is_text():
+        return
     click.secho(f"  {msg}", fg="green")
 
 
 def warn(msg: str) -> None:
+    if not is_text():
+        return
     click.secho(f"  WARNING: {msg}", fg="yellow", err=True)
 
 
@@ -92,8 +112,19 @@ class _AWSGroup(click.Group):
 
 @click.group(cls=_AWSGroup)
 @click.version_option(package_name="aws-bootstrap-g4dn")
-def main():
+@click.option(
+    "--output",
+    "-o",
+    type=click.Choice(["text", "json", "yaml", "table"], case_sensitive=False),
+    default="text",
+    show_default=True,
+    help="Output format.",
+)
+@click.pass_context
+def main(ctx, output):
     """Bootstrap AWS EC2 GPU instances for hybrid local-remote development."""
+    ctx.ensure_object(dict)
+    ctx.obj["output_format"] = OutputFormat(output)
 
 
 @main.command()
@@ -120,7 +151,21 @@ def main():
     help="Python version for the remote venv (e.g. 3.13, 3.14.2). Passed to uv during setup.",
 )
 @click.option("--ssh-port", default=22, show_default=True, type=int, help="SSH port on the remote instance.")
+@click.option(
+    "--ebs-storage",
+    default=None,
+    type=int,
+    help="Create and attach a new EBS data volume (size in GB, gp3). Mounted at /data.",
+)
+@click.option(
+    "--ebs-volume-id",
+    default=None,
+    type=str,
+    help="Attach an existing EBS volume by ID (e.g. vol-0abc123). Mounted at /data.",
+)
+@click.pass_context
 def launch(
+    ctx,
     instance_type,
     ami_filter,
     spot,
@@ -134,8 +179,13 @@ def launch(
     profile,
     python_version,
     ssh_port,
+    ebs_storage,
+    ebs_volume_id,
 ):
     """Launch a GPU-accelerated EC2 instance."""
+    if ebs_storage is not None and ebs_volume_id is not None:
+        raise CLIError("--ebs-storage and --ebs-volume-id are mutually exclusive.")
+
     config = LaunchConfig(
         instance_type=instance_type,
         spot=spot,
@@ -148,6 +198,8 @@ def launch(
         dry_run=dry_run,
         ssh_port=ssh_port,
         python_version=python_version,
+        ebs_storage=ebs_storage,
+        ebs_volume_id=ebs_volume_id,
     )
     if ami_filter:
         config.ami_filter = ami_filter
@@ -162,49 +214,78 @@ def launch(
     session = boto3.Session(profile_name=config.profile, region_name=config.region)
     ec2 = session.client("ec2")
 
+    has_ebs = config.ebs_storage is not None or config.ebs_volume_id is not None
+    total_steps = 7 if has_ebs else 6
+
     # Step 1: AMI lookup
-    step(1, 6, "Looking up AMI...")
+    step(1, total_steps, "Looking up AMI...")
     ami = get_latest_ami(ec2, config.ami_filter)
     info(f"Found: {ami['Name']}")
     val("AMI ID", ami["ImageId"])
 
     # Step 2: SSH key pair
-    step(2, 6, "Importing SSH key pair...")
+    step(2, total_steps, "Importing SSH key pair...")
     import_key_pair(ec2, config.key_name, config.key_path)
 
     # Step 3: Security group
-    step(3, 6, "Ensuring security group...")
+    step(3, total_steps, "Ensuring security group...")
     sg_id = ensure_security_group(ec2, config.security_group, config.tag_value, ssh_port=config.ssh_port)
 
     pricing = "spot" if config.spot else "on-demand"
 
     if config.dry_run:
-        click.echo()
-        click.secho("--- Dry Run Summary ---", bold=True, fg="yellow")
-        val("Instance type", config.instance_type)
-        val("AMI", f"{ami['ImageId']} ({ami['Name']})")
-        val("Pricing", pricing)
-        val("Key pair", config.key_name)
-        val("Security group", sg_id)
-        val("Volume", f"{config.volume_size} GB gp3")
-        val("Region", config.region)
-        val("Remote setup", "yes" if config.run_setup else "no")
-        if config.ssh_port != 22:
-            val("SSH port", str(config.ssh_port))
-        if config.python_version:
-            val("Python version", config.python_version)
-        click.echo()
-        click.secho("No resources launched (dry-run mode).", fg="yellow")
+        if is_text(ctx):
+            click.echo()
+            click.secho("--- Dry Run Summary ---", bold=True, fg="yellow")
+            val("Instance type", config.instance_type)
+            val("AMI", f"{ami['ImageId']} ({ami['Name']})")
+            val("Pricing", pricing)
+            val("Key pair", config.key_name)
+            val("Security group", sg_id)
+            val("Volume", f"{config.volume_size} GB gp3")
+            val("Region", config.region)
+            val("Remote setup", "yes" if config.run_setup else "no")
+            if config.ssh_port != 22:
+                val("SSH port", str(config.ssh_port))
+            if config.python_version:
+                val("Python version", config.python_version)
+            if config.ebs_storage:
+                val("EBS data volume", f"{config.ebs_storage} GB gp3 (new, mounted at {EBS_MOUNT_POINT})")
+            if config.ebs_volume_id:
+                val("EBS data volume", f"{config.ebs_volume_id} (existing, mounted at {EBS_MOUNT_POINT})")
+            click.echo()
+            click.secho("No resources launched (dry-run mode).", fg="yellow")
+        else:
+            result: dict = {
+                "dry_run": True,
+                "instance_type": config.instance_type,
+                "ami_id": ami["ImageId"],
+                "ami_name": ami["Name"],
+                "pricing": pricing,
+                "key_name": config.key_name,
+                "security_group": sg_id,
+                "volume_size_gb": config.volume_size,
+                "region": config.region,
+            }
+            if config.ssh_port != 22:
+                result["ssh_port"] = config.ssh_port
+            if config.python_version:
+                result["python_version"] = config.python_version
+            if config.ebs_storage:
+                result["ebs_storage_gb"] = config.ebs_storage
+            if config.ebs_volume_id:
+                result["ebs_volume_id"] = config.ebs_volume_id
+            emit(result, ctx=ctx)
         return
 
     # Step 4: Launch instance
-    step(4, 6, f"Launching {config.instance_type} instance ({pricing})...")
+    step(4, total_steps, f"Launching {config.instance_type} instance ({pricing})...")
     instance = launch_instance(ec2, config, ami["ImageId"], sg_id)
     instance_id = instance["InstanceId"]
     val("Instance ID", instance_id)
 
     # Step 5: Wait for ready
-    step(5, 6, "Waiting for instance to be ready...")
+    step(5, total_steps, "Waiting for instance to be ready...")
     instance = wait_instance_ready(ec2, instance_id)
     public_ip = instance.get("PublicIpAddress")
     if not public_ip:
@@ -213,9 +294,39 @@ def launch(
         return
 
     val("Public IP", public_ip)
+    az = instance["Placement"]["AvailabilityZone"]
+
+    # Step 5.5 (optional): EBS data volume
+    ebs_volume_attached = None
+    ebs_format = False
+    if has_ebs:
+        step(6, total_steps, "Setting up EBS data volume...")
+        if config.ebs_storage:
+            info(f"Creating {config.ebs_storage} GB gp3 volume in {az}...")
+            ebs_volume_attached = create_ebs_volume(ec2, config.ebs_storage, az, config.tag_value, instance_id)
+            val("Volume ID", ebs_volume_attached)
+            ebs_format = True
+        elif config.ebs_volume_id:
+            info(f"Validating volume {config.ebs_volume_id}...")
+            validate_ebs_volume(ec2, config.ebs_volume_id, az)
+            ebs_volume_attached = config.ebs_volume_id
+            # Tag the existing volume for discovery
+            ec2.create_tags(
+                Resources=[ebs_volume_attached],
+                Tags=[
+                    {"Key": "aws-bootstrap-instance", "Value": instance_id},
+                    {"Key": "created-by", "Value": config.tag_value},
+                ],
+            )
+            ebs_format = False
+
+        info(f"Attaching {ebs_volume_attached} to {instance_id}...")
+        attach_ebs_volume(ec2, ebs_volume_attached, instance_id)
+        success("EBS volume attached.")
 
-    # Step 6: SSH and remote setup
-    step(6, 6, "Waiting for SSH access...")
+    # SSH and remote setup step
+    ssh_step = 7 if has_ebs else 6
+    step(ssh_step, total_steps, "Waiting for SSH access...")
     private_key = private_key_path(config.key_path)
     if not wait_for_ssh(public_ip, config.ssh_user, config.key_path, port=config.ssh_port):
         warn("SSH did not become available within the timeout.")
@@ -238,6 +349,22 @@ def launch(
             else:
                 warn("Remote setup failed. Instance is still running.")
 
+    # Mount EBS volume via SSH (after setup so the instance is fully ready)
+    if ebs_volume_attached:
+        info(f"Mounting EBS volume at {EBS_MOUNT_POINT}...")
+        if mount_ebs_volume(
+            public_ip,
+            config.ssh_user,
+            config.key_path,
+            ebs_volume_attached,
+            mount_point=EBS_MOUNT_POINT,
+            format_volume=ebs_format,
+            port=config.ssh_port,
+        ):
+            success(f"EBS volume mounted at {EBS_MOUNT_POINT}.")
+        else:
+            warn(f"Failed to mount EBS volume at {EBS_MOUNT_POINT}. You may need to mount it manually.")
+
     # Add SSH config alias
     alias = add_ssh_host(
         instance_id=instance_id,
@@ -249,7 +376,30 @@ def launch(
     )
     success(f"Added SSH config alias: {alias}")
 
-    # Print connection info
+    # Structured output for non-text modes
+    if not is_text(ctx):
+        result_data: dict = {
+            "instance_id": instance_id,
+            "public_ip": public_ip,
+            "instance_type": config.instance_type,
+            "availability_zone": az,
+            "ami_id": ami["ImageId"],
+            "pricing": pricing,
+            "region": config.region,
+            "ssh_alias": alias,
+        }
+        if ebs_volume_attached:
+            ebs_info: dict = {
+                "volume_id": ebs_volume_attached,
+                "mount_point": EBS_MOUNT_POINT,
+            }
+            if config.ebs_storage:
+                ebs_info["size_gb"] = config.ebs_storage
+            result_data["ebs_volume"] = ebs_info
+        emit(result_data, ctx=ctx)
+        return
+
+    # Print connection info (text mode)
     click.echo()
     click.secho("=" * 60, fg="green")
     click.secho("  Instance ready!", bold=True, fg="green")
@@ -260,6 +410,12 @@ def launch(
     val("Instance", config.instance_type)
     val("Pricing", pricing)
     val("SSH alias", alias)
+    if ebs_volume_attached:
+        if config.ebs_storage:
+            ebs_label = f"{ebs_volume_attached} ({config.ebs_storage} GB, {EBS_MOUNT_POINT})"
+        else:
+            ebs_label = f"{ebs_volume_attached} ({EBS_MOUNT_POINT})"
+        val("EBS data volume", ebs_label)
 
     port_flag = f" -p {config.ssh_port}" if config.ssh_port != 22 else ""
 
@@ -304,44 +460,66 @@ def launch(
     show_default=True,
     help="Show connection commands (SSH, Jupyter, VSCode) for each running instance.",
 )
-def status(region, profile, gpu, instructions):
+@click.pass_context
+def status(ctx, region, profile, gpu, instructions):
     """Show running instances created by aws-bootstrap."""
     session = boto3.Session(profile_name=profile, region_name=region)
     ec2 = session.client("ec2")
 
     instances = find_tagged_instances(ec2, "aws-bootstrap-g4dn")
     if not instances:
-        click.secho("No active aws-bootstrap instances found.", fg="yellow")
+        if is_text(ctx):
+            click.secho("No active aws-bootstrap instances found.", fg="yellow")
+        else:
+            emit({"instances": []}, ctx=ctx)
         return
 
     ssh_hosts = list_ssh_hosts()
 
-    click.secho(f"\n  Found {len(instances)} instance(s):\n", bold=True, fg="cyan")
-    if gpu:
-        click.echo("  " + click.style("Querying GPU info via SSH...", dim=True))
-        click.echo()
+    if is_text(ctx):
+        click.secho(f"\n  Found {len(instances)} instance(s):\n", bold=True, fg="cyan")
+        if gpu:
+            click.echo("  " + click.style("Querying GPU info via SSH...", dim=True))
+            click.echo()
+
+    structured_instances = []
 
     for inst in instances:
         state = inst["State"]
-        state_color = {
-            "running": "green",
-            "pending": "yellow",
-            "stopping": "yellow",
-            "stopped": "red",
-            "shutting-down": "red",
-        }.get(state, "white")
         alias = ssh_hosts.get(inst["InstanceId"])
-        alias_str = f" ({alias})" if alias else ""
-        click.echo(
-            "  "
-            + click.style(inst["InstanceId"], fg="bright_white")
-            + click.style(alias_str, fg="cyan")
-            + "  "
-            + click.style(state, fg=state_color)
-        )
-        val("    Type", inst["InstanceType"])
-        if inst["PublicIp"]:
-            val("    IP", inst["PublicIp"])
+
+        # Text mode: inline display
+        if is_text(ctx):
+            state_color = {
+                "running": "green",
+                "pending": "yellow",
+                "stopping": "yellow",
+                "stopped": "red",
+                "shutting-down": "red",
+            }.get(state, "white")
+            alias_str = f" ({alias})" if alias else ""
+            click.echo(
+                "  "
+                + click.style(inst["InstanceId"], fg="bright_white")
+                + click.style(alias_str, fg="cyan")
+                + "  "
+                + click.style(state, fg=state_color)
+            )
+            val("    Type", inst["InstanceType"])
+            if inst["PublicIp"]:
+                val("    IP", inst["PublicIp"])
+
+        # Build structured record
+        inst_data: dict = {
+            "instance_id": inst["InstanceId"],
+            "state": state,
+            "instance_type": inst["InstanceType"],
+            "public_ip": inst["PublicIp"] or None,
+            "ssh_alias": alias,
+            "lifecycle": inst["Lifecycle"],
+            "availability_zone": inst["AvailabilityZone"],
+            "launch_time": inst["LaunchTime"],
+        }
 
         # Look up SSH config details once (used by --gpu and --with-instructions)
         details = None
@@ -359,45 +537,81 @@ def status(region, profile, gpu, instructions):
                     Path("~/.ssh/id_ed25519").expanduser(),
                 )
             if gpu_info:
-                val("    GPU", f"{gpu_info.gpu_name} ({gpu_info.architecture})")
-                if gpu_info.cuda_toolkit_version:
-                    cuda_str = gpu_info.cuda_toolkit_version
-                    if gpu_info.cuda_driver_version != gpu_info.cuda_toolkit_version:
-                        cuda_str += f" (driver supports up to {gpu_info.cuda_driver_version})"
-                else:
-                    cuda_str = f"{gpu_info.cuda_driver_version} (driver max, toolkit unknown)"
-                val("    CUDA", cuda_str)
-                val("    Driver", gpu_info.driver_version)
+                if is_text(ctx):
+                    val("    GPU", f"{gpu_info.gpu_name} ({gpu_info.architecture})")
+                    if gpu_info.cuda_toolkit_version:
+                        cuda_str = gpu_info.cuda_toolkit_version
+                        if gpu_info.cuda_driver_version != gpu_info.cuda_toolkit_version:
+                            cuda_str += f" (driver supports up to {gpu_info.cuda_driver_version})"
+                    else:
+                        cuda_str = f"{gpu_info.cuda_driver_version} (driver max, toolkit unknown)"
+                    val("    CUDA", cuda_str)
+                    val("    Driver", gpu_info.driver_version)
+                inst_data["gpu"] = {
+                    "name": gpu_info.gpu_name,
+                    "architecture": gpu_info.architecture,
+                    "cuda_toolkit": gpu_info.cuda_toolkit_version,
+                    "cuda_driver_max": gpu_info.cuda_driver_version,
+                    "driver": gpu_info.driver_version,
+                }
             else:
-                click.echo("    GPU: " + click.style("unavailable", dim=True))
+                if is_text(ctx):
+                    click.echo("    GPU: " + click.style("unavailable", dim=True))
+
+        # EBS data volumes
+        ebs_volumes = find_ebs_volumes_for_instance(ec2, inst["InstanceId"], "aws-bootstrap-g4dn")
+        if ebs_volumes:
+            if is_text(ctx):
+                for vol in ebs_volumes:
+                    vol_state = f", {vol['State']}" if vol["State"] != "in-use" else ""
+                    val("    EBS", f"{vol['VolumeId']} ({vol['Size']} GB, {EBS_MOUNT_POINT}{vol_state})")
+            inst_data["ebs_volumes"] = [
+                {
+                    "volume_id": vol["VolumeId"],
+                    "size_gb": vol["Size"],
+                    "mount_point": EBS_MOUNT_POINT,
+                    "state": vol["State"],
+                }
+                for vol in ebs_volumes
+            ]
 
         lifecycle = inst["Lifecycle"]
         is_spot = lifecycle == "spot"
+        spot_price = None
 
         if is_spot:
             spot_price = get_spot_price(ec2, inst["InstanceType"], inst["AvailabilityZone"])
+            if is_text(ctx):
+                if spot_price is not None:
+                    val("    Pricing", f"spot (${spot_price:.4f}/hr)")
+                else:
+                    val("    Pricing", "spot")
             if spot_price is not None:
-                val("    Pricing", f"spot (${spot_price:.4f}/hr)")
-            else:
-                val("    Pricing", "spot")
+                inst_data["spot_price_per_hour"] = spot_price
         else:
-            val("    Pricing", "on-demand")
+            if is_text(ctx):
+                val("    Pricing", "on-demand")
 
         if state == "running" and is_spot:
             uptime = datetime.now(UTC) - inst["LaunchTime"]
             total_seconds = int(uptime.total_seconds())
-            hours, remainder = divmod(total_seconds, 3600)
-            minutes = remainder // 60
-            val("    Uptime", f"{hours}h {minutes:02d}m")
+            inst_data["uptime_seconds"] = total_seconds
+            if is_text(ctx):
+                hours, remainder = divmod(total_seconds, 3600)
+                minutes = remainder // 60
+                val("    Uptime", f"{hours}h {minutes:02d}m")
             if spot_price is not None:
                 uptime_hours = uptime.total_seconds() / 3600
                 est_cost = uptime_hours * spot_price
-                val("    Est. cost", f"~${est_cost:.4f}")
+                inst_data["estimated_cost"] = round(est_cost, 4)
+                if is_text(ctx):
+                    val("    Est. cost", f"~${est_cost:.4f}")
 
-        val("    Launched", str(inst["LaunchTime"]))
+        if is_text(ctx):
+            val("    Launched", str(inst["LaunchTime"]))
 
         # Connection instructions (opt-in, only for running instances with a public IP and alias)
-        if instructions and state == "running" and inst["PublicIp"] and alias:
+        if is_text(ctx) and instructions and state == "running" and inst["PublicIp"] and alias:
             user = details.user if details else "ubuntu"
             port = details.port if details else 22
             port_flag = f" -p {port}" if port != 22 else ""
@@ -418,6 +632,24 @@ def status(region, profile, gpu, instructions):
             click.secho("    GPU Benchmark:", fg="cyan")
             click.secho(f"      ssh {alias} 'python ~/gpu_benchmark.py'", bold=True)
 
+        structured_instances.append(inst_data)
+
+    if not is_text(ctx):
+        emit(
+            {"instances": structured_instances},
+            headers={
+                "instance_id": "Instance ID",
+                "state": "State",
+                "instance_type": "Type",
+                "public_ip": "IP",
+                "ssh_alias": "Alias",
+                "lifecycle": "Pricing",
+                "uptime_seconds": "Uptime (s)",
+            },
+            ctx=ctx,
+        )
+        return
+
     click.echo()
     first_id = instances[0]["InstanceId"]
     first_ref = ssh_hosts.get(first_id, first_id)
@@ -429,8 +661,10 @@ def status(region, profile, gpu, instructions):
 @click.option("--region", default="us-west-2", show_default=True, help="AWS region.")
 @click.option("--profile", default=None, help="AWS profile override.")
 @click.option("--yes", "-y", is_flag=True, default=False, help="Skip confirmation prompt.")
+@click.option("--keep-ebs", is_flag=True, default=False, help="Preserve EBS data volumes instead of deleting them.")
 @click.argument("instance_ids", nargs=-1, metavar="[INSTANCE_ID_OR_ALIAS]...")
-def terminate(region, profile, yes, instance_ids):
+@click.pass_context
+def terminate(ctx, region, profile, yes, keep_ebs, instance_ids):
     """Terminate instances created by aws-bootstrap.
 
     Pass specific instance IDs or SSH aliases (e.g. aws-gpu1) to terminate,
@@ -439,6 +673,10 @@ def terminate(region, profile, yes, instance_ids):
     session = boto3.Session(profile_name=profile, region_name=region)
     ec2 = session.client("ec2")
 
+    # In structured output modes, require --yes (prompts would corrupt output)
+    if not is_text(ctx) and not yes:
+        raise CLIError("--yes is required when using structured output (--output json/yaml/table).")
+
     if instance_ids:
         targets = []
         for value in instance_ids:
@@ -454,13 +692,17 @@ def terminate(region, profile, yes, instance_ids):
     else:
         instances = find_tagged_instances(ec2, "aws-bootstrap-g4dn")
         if not instances:
-            click.secho("No active aws-bootstrap instances found.", fg="yellow")
+            if is_text(ctx):
+                click.secho("No active aws-bootstrap instances found.", fg="yellow")
+            else:
+                emit({"terminated": []}, ctx=ctx)
             return
         targets = [inst["InstanceId"] for inst in instances]
-        click.secho(f"\n  Found {len(targets)} instance(s) to terminate:\n", bold=True, fg="cyan")
-        for inst in instances:
-            iid = click.style(inst["InstanceId"], fg="bright_white")
-            click.echo(f"  {iid}  {inst['State']}  {inst['InstanceType']}")
+        if is_text(ctx):
+            click.secho(f"\n  Found {len(targets)} instance(s) to terminate:\n", bold=True, fg="cyan")
+            for inst in instances:
+                iid = click.style(inst["InstanceId"], fg="bright_white")
+                click.echo(f"  {iid}  {inst['State']}  {inst['InstanceType']}")
 
     if not yes:
         click.echo()
@@ -468,21 +710,143 @@ def terminate(region, profile, yes, instance_ids):
             click.secho("  Cancelled.", fg="yellow")
             return
 
+    # Discover EBS volumes before termination (while instances still exist)
+    ebs_by_instance: dict[str, list[dict]] = {}
+    for target in targets:
+        volumes = find_ebs_volumes_for_instance(ec2, target, "aws-bootstrap-g4dn")
+        if volumes:
+            ebs_by_instance[target] = volumes
+
     changes = terminate_tagged_instances(ec2, targets)
-    click.echo()
+
+    terminated_results = []
+
+    if is_text(ctx):
+        click.echo()
     for change in changes:
         prev = change["PreviousState"]["Name"]
         curr = change["CurrentState"]["Name"]
-        click.echo(
-            "  " + click.style(change["InstanceId"], fg="bright_white") + f"  {prev} -> " + click.style(curr, fg="red")
-        )
-        removed_alias = remove_ssh_host(change["InstanceId"])
+        iid = change["InstanceId"]
+        if is_text(ctx):
+            click.echo("  " + click.style(iid, fg="bright_white") + f"  {prev} -> " + click.style(curr, fg="red"))
+        removed_alias = remove_ssh_host(iid)
         if removed_alias:
             info(f"Removed SSH config alias: {removed_alias}")
+
+        change_data: dict = {
+            "instance_id": iid,
+            "previous_state": prev,
+            "current_state": curr,
+        }
+        if removed_alias:
+            change_data["ssh_alias_removed"] = removed_alias
+        terminated_results.append(change_data)
+
+    # Handle EBS volume cleanup
+    for _iid, volumes in ebs_by_instance.items():
+        for vol in volumes:
+            vid = vol["VolumeId"]
+            if keep_ebs:
+                if is_text(ctx):
+                    click.echo()
+                info(f"Preserving EBS volume: {vid} ({vol['Size']} GB)")
+                info(f"Reattach with: aws-bootstrap launch --ebs-volume-id {vid}")
+            else:
+                if is_text(ctx):
+                    click.echo()
+                info(f"Waiting for EBS volume {vid} to detach...")
+                try:
+                    waiter = ec2.get_waiter("volume_available")
+                    waiter.wait(VolumeIds=[vid], WaiterConfig={"Delay": 10, "MaxAttempts": 30})
+                    delete_ebs_volume(ec2, vid)
+                    success(f"Deleted EBS volume: {vid}")
+                    # Record deleted volume in the corresponding terminated result
+                    for tr in terminated_results:
+                        if tr["instance_id"] == _iid:
+                            tr.setdefault("ebs_volumes_deleted", []).append(vid)
+                except Exception as e:
+                    warn(f"Failed to delete EBS volume {vid}: {e}")
+
+    if not is_text(ctx):
+        emit({"terminated": terminated_results}, ctx=ctx)
+        return
+
     click.echo()
     success(f"Terminated {len(changes)} instance(s).")
 
 
+@main.command()
+@click.option("--dry-run", is_flag=True, default=False, help="Show what would be removed without removing.")
+@click.option("--yes", "-y", is_flag=True, default=False, help="Skip confirmation prompt.")
+@click.option("--region", default="us-west-2", show_default=True, help="AWS region.")
+@click.option("--profile", default=None, help="AWS profile override.")
+@click.pass_context
+def cleanup(ctx, dry_run, yes, region, profile):
+    """Remove stale SSH config entries for terminated instances."""
+    session = boto3.Session(profile_name=profile, region_name=region)
+    ec2 = session.client("ec2")
+
+    # In structured output modes, require --yes for non-dry-run (prompts would corrupt output)
+    if not is_text(ctx) and not yes and not dry_run:
+        raise CLIError("--yes is required when using structured output (--output json/yaml/table).")
+
+    live_instances = find_tagged_instances(ec2, "aws-bootstrap-g4dn")
+    live_ids = {inst["InstanceId"] for inst in live_instances}
+
+    stale = find_stale_ssh_hosts(live_ids)
+    if not stale:
+        if is_text(ctx):
+            click.secho("No stale SSH config entries found.", fg="green")
+        else:
+            result_key = "stale" if dry_run else "cleaned"
+            emit({result_key: []}, ctx=ctx)
+        return
+
+    if is_text(ctx):
+        click.secho(f"\n  Found {len(stale)} stale SSH config entry(ies):\n", bold=True, fg="cyan")
+        for iid, alias in stale:
+            click.echo("  " + click.style(alias, fg="bright_white") + f"  ({iid})")
+
+    if dry_run:
+        if is_text(ctx):
+            click.echo()
+            for iid, alias in stale:
+                info(f"Would remove {alias} ({iid})")
+        else:
+            emit(
+                {
+                    "stale": [{"instance_id": iid, "alias": alias} for iid, alias in stale],
+                    "dry_run": True,
+                },
+                ctx=ctx,
+            )
+        return
+
+    if not yes:
+        click.echo()
+        if not click.confirm(f"  Remove {len(stale)} stale entry(ies)?"):
+            click.secho("  Cancelled.", fg="yellow")
+            return
+
+    results = cleanup_stale_ssh_hosts(live_ids)
+
+    if not is_text(ctx):
+        emit(
+            {
+                "cleaned": [{"instance_id": r.instance_id, "alias": r.alias, "removed": r.removed} for r in results],
+            },
+            ctx=ctx,
+        )
+        return
+
+    click.echo()
+    for r in results:
+        success(f"Removed {r.alias} ({r.instance_id})")
+
+    click.echo()
+    success(f"Cleaned up {len(results)} stale entry(ies).")
+
+
 # ---------------------------------------------------------------------------
 # list command group
 # ---------------------------------------------------------------------------
@@ -499,14 +863,40 @@ def list_cmd():
 @click.option("--prefix", default="g4dn", show_default=True, help="Instance type family prefix to filter on.")
 @click.option("--region", default="us-west-2", show_default=True, help="AWS region.")
 @click.option("--profile", default=None, help="AWS profile override.")
-def list_instance_types_cmd(prefix, region, profile):
+@click.pass_context
+def list_instance_types_cmd(ctx, prefix, region, profile):
     """List EC2 instance types matching a family prefix (e.g. g4dn, p3, g5)."""
     session = boto3.Session(profile_name=profile, region_name=region)
     ec2 = session.client("ec2")
 
     types = list_instance_types(ec2, prefix)
     if not types:
-        click.secho(f"No instance types found matching '{prefix}.*'", fg="yellow")
+        if is_text(ctx):
+            click.secho(f"No instance types found matching '{prefix}.*'", fg="yellow")
+        else:
+            emit([], ctx=ctx)
+        return
+
+    if not is_text(ctx):
+        structured = [
+            {
+                "instance_type": t["InstanceType"],
+                "vcpus": t["VCpuCount"],
+                "memory_mib": t["MemoryMiB"],
+                "gpu": t["GpuSummary"] or None,
+            }
+            for t in types
+        ]
+        emit(
+            structured,
+            headers={
+                "instance_type": "Instance Type",
+                "vcpus": "vCPUs",
+                "memory_mib": "Memory (MiB)",
+                "gpu": "GPU",
+            },
+            ctx=ctx,
+        )
         return
 
     click.secho(f"\n  {len(types)} instance type(s) matching '{prefix}.*':\n", bold=True, fg="cyan")
@@ -518,8 +908,8 @@ def list_instance_types_cmd(prefix, region, profile):
     click.echo("  " + "-" * 72)
 
     for t in types:
-        gpu = t["GpuSummary"] or "-"
-        click.echo(f"  {t['InstanceType']:<24}{t['VCpuCount']:>6}{t['MemoryMiB']:>14}  {gpu}")
+        gpu_str = t["GpuSummary"] or "-"
+        click.echo(f"  {t['InstanceType']:<24}{t['VCpuCount']:>6}{t['MemoryMiB']:>14}  {gpu_str}")
 
     click.echo()
 
@@ -528,14 +918,40 @@ def list_instance_types_cmd(prefix, region, profile):
 @click.option("--filter", "ami_filter", default=DEFAULT_AMI_PREFIX, show_default=True, help="AMI name pattern.")
 @click.option("--region", default="us-west-2", show_default=True, help="AWS region.")
 @click.option("--profile", default=None, help="AWS profile override.")
-def list_amis_cmd(ami_filter, region, profile):
+@click.pass_context
+def list_amis_cmd(ctx, ami_filter, region, profile):
     """List available AMIs matching a name pattern."""
     session = boto3.Session(profile_name=profile, region_name=region)
     ec2 = session.client("ec2")
 
     amis = list_amis(ec2, ami_filter)
     if not amis:
-        click.secho(f"No AMIs found matching '{ami_filter}'", fg="yellow")
+        if is_text(ctx):
+            click.secho(f"No AMIs found matching '{ami_filter}'", fg="yellow")
+        else:
+            emit([], ctx=ctx)
+        return
+
+    if not is_text(ctx):
+        structured = [
+            {
+                "image_id": ami["ImageId"],
+                "name": ami["Name"],
+                "creation_date": ami["CreationDate"][:10],
+                "architecture": ami["Architecture"],
+            }
+            for ami in amis
+        ]
+        emit(
+            structured,
+            headers={
+                "image_id": "Image ID",
+                "name": "Name",
+                "creation_date": "Created",
+                "architecture": "Arch",
+            },
+            ctx=ctx,
+        )
         return
 
     click.secho(f"\n  {len(amis)} AMI(s) matching '{ami_filter}' (newest first):\n", bold=True, fg="cyan")