aws-bootstrap-g4dn 0.5.0__tar.gz → 0.6.0__tar.gz

{aws_bootstrap_g4dn-0.5.0 → aws_bootstrap_g4dn-0.6.0}/CLAUDE.md

@@ -32,9 +32,9 @@ aws_bootstrap/
     __init__.py          # Package init
     cli.py               # Click CLI entry point (launch, status, terminate commands)
     config.py            # LaunchConfig dataclass with defaults
-    ec2.py               # AMI lookup, security group, instance launch/find/terminate, polling, spot pricing
+    ec2.py               # AMI lookup, security group, instance launch/find/terminate, polling, spot pricing, EBS volume ops
     gpu.py               # GPU architecture mapping and GpuInfo dataclass
-    ssh.py               # SSH key pair import, SSH readiness check, remote setup, ~/.ssh/config management, GPU queries
+    ssh.py               # SSH key pair import, SSH readiness check, remote setup, ~/.ssh/config management, GPU queries, EBS mount
     resources/           # Non-Python artifacts SCP'd to remote instances
         __init__.py
         gpu_benchmark.py       # GPU throughput benchmark (CNN + Transformer), copied to ~/gpu_benchmark.py on instance
@@ -51,6 +51,8 @@ aws_bootstrap/
         test_gpu.py
         test_ssh_config.py
         test_ssh_gpu.py
+        test_ebs.py
+        test_ssh_ebs.py
 docs/
     nsight-remote-profiling.md # Nsight Compute, Nsight Systems, and Nsight VSCE remote profiling guide
     spot-request-lifecycle.md  # Research notes on spot request cleanup
@@ -60,9 +62,10 @@ Entry point: `aws-bootstrap = "aws_bootstrap.cli:main"` (installed via `uv sync`
 
 ## CLI Commands
 
-- **`launch`** — provisions an EC2 instance (spot by default, falls back to on-demand on capacity errors); adds SSH config alias (e.g. `aws-gpu1`) to `~/.ssh/config`; `--python-version` controls which Python `uv` installs in the remote venv; `--ssh-port` overrides the default SSH port (22) for security group ingress, connection checks, and SSH config
-- **`status`** — lists all non-terminated instances (including `shutting-down`) with type, IP, SSH alias, pricing (spot price/hr or on-demand), uptime, and estimated cost for running spot instances; `--gpu` flag queries GPU info via SSH, reporting both CUDA toolkit version (from `nvcc`) and driver-supported max (from `nvidia-smi`); `--instructions` (default: on) prints connection commands (SSH, Jupyter tunnel, VSCode Remote SSH, GPU benchmark) for each running instance; suppress with `--no-instructions`
-- **`terminate`** — terminates instances by ID or SSH alias (e.g. `aws-gpu1`, resolved via `~/.ssh/config`), or all aws-bootstrap instances in the region if no arguments given; removes SSH config aliases
+- **`launch`** — provisions an EC2 instance (spot by default, falls back to on-demand on capacity errors); adds SSH config alias (e.g. `aws-gpu1`) to `~/.ssh/config`; `--python-version` controls which Python `uv` installs in the remote venv; `--ssh-port` overrides the default SSH port (22) for security group ingress, connection checks, and SSH config; `--ebs-storage SIZE` creates and attaches a new gp3 EBS data volume (mounted at `/data`); `--ebs-volume-id ID` attaches an existing EBS volume (mutually exclusive with `--ebs-storage`)
+- **`status`** — lists all non-terminated instances (including `shutting-down`) with type, IP, SSH alias, EBS data volumes, pricing (spot price/hr or on-demand), uptime, and estimated cost for running spot instances; `--gpu` flag queries GPU info via SSH, reporting both CUDA toolkit version (from `nvcc`) and driver-supported max (from `nvidia-smi`); `--instructions` (default: on) prints connection commands (SSH, Jupyter tunnel, VSCode Remote SSH, GPU benchmark) for each running instance; suppress with `--no-instructions`
+- **`terminate`** — terminates instances by ID or SSH alias (e.g. `aws-gpu1`, resolved via `~/.ssh/config`), or all aws-bootstrap instances in the region if no arguments given; removes SSH config aliases; deletes associated EBS data volumes by default; `--keep-ebs` preserves volumes and prints reattach commands
+- **`cleanup`** — removes stale `~/.ssh/config` entries for terminated/non-existent instances; compares managed SSH config blocks against live EC2 instances; `--dry-run` previews removals without modifying config; `--yes` skips the confirmation prompt
 - **`list instance-types`** — lists EC2 instance types matching a family prefix (default: `g4dn`), showing vCPUs, memory, and GPU info
 - **`list amis`** — lists available AMIs matching a name pattern (default: Deep Learning Base OSS Nvidia Driver GPU AMIs), sorted newest-first
 
@@ -112,6 +115,34 @@ The `KNOWN_CUDA_TAGS` array in `remote_setup.sh` lists the CUDA wheel tags publi
 
 `resources/gpu_benchmark.py` is uploaded to `~/gpu_benchmark.py` on the remote instance during setup. It benchmarks GPU throughput with two modes: CNN on MNIST and a GPT-style Transformer on synthetic data. It reports samples/sec, batch times, and peak GPU memory. Supports `--precision` (fp32/fp16/bf16/tf32), `--diagnose` for CUDA smoke tests, and separate `--transformer-batch-size` (default 32, T4-safe). Dependencies (`torch`, `torchvision`, `tqdm`) are already installed by the setup script.
 
+## EBS Data Volumes
+
+The `--ebs-storage` and `--ebs-volume-id` options on `launch` create or attach persistent gp3 EBS volumes mounted at `/data`. The implementation spans three modules:
+
+- **`ec2.py`** — Volume lifecycle: `create_ebs_volume`, `validate_ebs_volume`, `attach_ebs_volume`, `detach_ebs_volume`, `delete_ebs_volume`, `find_ebs_volumes_for_instance`. Constants `EBS_DEVICE_NAME` (`/dev/sdf`) and `EBS_MOUNT_POINT` (`/data`).
+- **`ssh.py`** — `mount_ebs_volume()` SSHs to the instance and runs a shell script that detects the device, optionally formats it, mounts it, and adds an fstab entry.
+- **`cli.py`** — Orchestrates the flow: create/validate → attach → wait for SSH → mount. Mount failures are non-fatal (warn and continue).
+
+### Tagging strategy
+
+Volumes are tagged for discovery by `status` and `terminate`:
+
+| Tag | Value | Purpose |
+|-----|-------|---------|
+| `created-by` | `aws-bootstrap-g4dn` | Standard tool-managed resource tag |
+| `Name` | `aws-bootstrap-data-{instance_id}` | Human-readable in AWS console |
+| `aws-bootstrap-instance` | `i-xxxxxxxxx` | Links volume to instance for `find_ebs_volumes_for_instance` |
+
+### NVMe device detection
+
+On Nitro instances (g4dn), `/dev/sdf` is remapped to `/dev/nvmeXn1`. The mount script detects the correct device by matching the volume ID serial number via `lsblk -o NAME,SERIAL -dpn`, with fallbacks to `/dev/nvme1n1`, `/dev/xvdf`, `/dev/sdf`.
+
+### Spot interruption and terminate cleanup
+
+Non-root EBS volumes attached via API have `DeleteOnTermination=False` by default. This means data volumes **survive spot interruptions** — when AWS reclaims the instance, the volume detaches and becomes `available`, preserving all data. The user can reattach it to a new instance with `--ebs-volume-id`.
+
+The `terminate` command discovers volumes via `find_ebs_volumes_for_instance`, waits for them to detach (becomes `available`), then deletes them. `--keep-ebs` skips deletion and prints the volume ID with a reattach command.
+
 ## Versioning & Publishing
 
 Version is derived automatically from git tags via **setuptools-scm** — no hardcoded version string in the codebase.

{aws_bootstrap_g4dn-0.5.0 → aws_bootstrap_g4dn-0.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aws-bootstrap-g4dn
-Version: 0.5.0
+Version: 0.6.0
 Summary: Bootstrap AWS EC2 GPU instances for hybrid local-remote development
 Author: Adam Ever-Hadani
 License-Expression: MIT
@@ -44,7 +44,8 @@ ssh aws-gpu1                  # You're in, venv activated, PyTorch works
 | 📊 | **GPU benchmark included** | CNN (MNIST) + Transformer benchmarks with FP16/FP32/BF16 precision and tqdm progress |
 | 📓 | **Jupyter ready** | Lab server auto-starts as a systemd service on port 8888 — just SSH tunnel and open |
 | 🖥️ | **`status --gpu`** | Shows CUDA toolkit version, driver max, GPU architecture, spot pricing, uptime, and estimated cost |
-| 🗑️ | **Clean terminate** | Stops instances, removes SSH aliases, shows shutting-down state until fully gone |
+| 💾 | **EBS data volumes** | Attach persistent storage at `/data` — survives spot interruptions and termination, reattach to new instances |
+| 🗑️ | **Clean terminate** | Stops instances, removes SSH aliases, cleans up EBS volumes (or preserves with `--keep-ebs`) |
 
 ### 🎯 Target Workflows
 
@@ -132,16 +133,24 @@ aws-bootstrap launch --python-version 3.13
 # Use a non-default SSH port
 aws-bootstrap launch --ssh-port 2222
 
+# Attach a persistent EBS data volume (96 GB gp3, mounted at /data)
+aws-bootstrap launch --ebs-storage 96
+
+# Reattach an existing EBS volume from a previous instance
+aws-bootstrap launch --ebs-volume-id vol-0abc123def456
+
 # Use a specific AWS profile
 aws-bootstrap launch --profile my-aws-profile
 ```
 
 After launch, the CLI:
 
-1. **Adds an SSH alias** (e.g. `aws-gpu1`) to `~/.ssh/config`
-2. **Runs remote setup** — installs utilities, creates a Python venv, installs CUDA-matched PyTorch, sets up Jupyter
-3. **Runs a CUDA smoke test** — verifies `torch.cuda.is_available()` and runs a quick GPU matmul
-4. **Prints connection commands** — SSH, Jupyter tunnel, GPU benchmark, and terminate
+1. **Creates/attaches EBS volume** (if `--ebs-storage` or `--ebs-volume-id` was specified)
+2. **Adds an SSH alias** (e.g. `aws-gpu1`) to `~/.ssh/config`
+3. **Runs remote setup** — installs utilities, creates a Python venv, installs CUDA-matched PyTorch, sets up Jupyter
+4. **Mounts EBS volume** at `/data` (if applicable — formats new volumes, mounts existing ones as-is)
+5. **Runs a CUDA smoke test** — verifies `torch.cuda.is_available()` and runs a quick GPU matmul
+6. **Prints connection commands** — SSH, Jupyter tunnel, GPU benchmark, and terminate
 
 ```bash
 ssh aws-gpu1                  # venv auto-activates on login
@@ -154,7 +163,7 @@ The setup script runs automatically on the instance after SSH becomes available:
 | Step | What |
 |------|------|
 | **GPU verify** | Confirms `nvidia-smi` and `nvcc` are working |
-| **Utilities** | Installs `htop`, `tmux`, `tree`, `jq` |
+| **Utilities** | Installs `htop`, `tmux`, `tree`, `jq`, `ffmpeg` |
 | **Python venv** | Creates `~/venv` with `uv`, auto-activates in `~/.bashrc`. Use `--python-version` to pin a specific Python (e.g. `3.13`) |
 | **CUDA-aware PyTorch** | Detects CUDA toolkit version → installs PyTorch from the matching `cu{TAG}` wheel index |
 | **CUDA smoke test** | Runs `torch.cuda.is_available()` + GPU matmul to verify the stack |
@@ -261,6 +270,9 @@ aws-bootstrap status --region us-east-1
 # Terminate all aws-bootstrap instances (with confirmation prompt)
 aws-bootstrap terminate
 
+# Terminate but preserve EBS data volumes for reuse
+aws-bootstrap terminate --keep-ebs
+
 # Terminate by SSH alias (resolved via ~/.ssh/config)
 aws-bootstrap terminate aws-gpu1
 
@@ -272,6 +284,15 @@ aws-bootstrap terminate aws-gpu1 i-def456
 
 # Skip confirmation prompt
 aws-bootstrap terminate --yes
+
+# Remove stale SSH config entries for terminated instances
+aws-bootstrap cleanup
+
+# Preview what would be removed without modifying config
+aws-bootstrap cleanup --dry-run
+
+# Skip confirmation prompt
+aws-bootstrap cleanup --yes
 ```
 
 `status --gpu` reports both the **installed CUDA toolkit** version (from `nvcc`) and the **maximum CUDA version supported by the driver** (from `nvidia-smi`), so you can see at a glance whether they match:
@@ -282,6 +303,31 @@ CUDA: 12.8 (driver supports up to 13.0)
 
 SSH aliases are managed automatically — they're created on `launch`, shown in `status`, and cleaned up on `terminate`. Aliases use sequential numbering (`aws-gpu1`, `aws-gpu2`, etc.) and never reuse numbers from previous instances. You can use aliases anywhere you'd use an instance ID, e.g. `aws-bootstrap terminate aws-gpu1`.
 
+## EBS Data Volumes
+
+Attach persistent EBS storage to keep datasets and model checkpoints across instance lifecycles. Volumes are mounted at `/data` and persist independently of the instance.
+
+```bash
+# Create a new 96 GB gp3 volume, formatted and mounted at /data
+aws-bootstrap launch --ebs-storage 96
+
+# After terminating with --keep-ebs, reattach the same volume to a new instance
+aws-bootstrap terminate --keep-ebs
+# Output: Preserving EBS volume: vol-0abc123...
+#         Reattach with: aws-bootstrap launch --ebs-volume-id vol-0abc123...
+
+aws-bootstrap launch --ebs-volume-id vol-0abc123def456
+```
+
+Key behaviors:
+- `--ebs-storage` and `--ebs-volume-id` are mutually exclusive
+- New volumes are formatted as ext4; existing volumes are mounted as-is
+- Volumes are tagged for automatic discovery by `status` and `terminate`
+- `terminate` deletes data volumes by default; use `--keep-ebs` to preserve them
+- **Spot-safe** — data volumes survive spot interruptions. If AWS reclaims your instance, the volume detaches automatically and can be reattached to a new instance with `--ebs-volume-id`
+- EBS volumes must be in the same availability zone as the instance
+- Mount failures are non-fatal — the instance remains usable
+
 ## EC2 vCPU Quotas
 
 AWS accounts have [service quotas](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-resource-limits.html) that limit how many vCPUs you can run per instance family. New or lightly-used accounts often have a **default quota of 0 vCPUs** for GPU instance families (G and VT), which will cause errors on launch:

{aws_bootstrap_g4dn-0.5.0 → aws_bootstrap_g4dn-0.6.0}/README.md

@@ -25,7 +25,8 @@ ssh aws-gpu1                  # You're in, venv activated, PyTorch works
 | 📊 | **GPU benchmark included** | CNN (MNIST) + Transformer benchmarks with FP16/FP32/BF16 precision and tqdm progress |
 | 📓 | **Jupyter ready** | Lab server auto-starts as a systemd service on port 8888 — just SSH tunnel and open |
 | 🖥️ | **`status --gpu`** | Shows CUDA toolkit version, driver max, GPU architecture, spot pricing, uptime, and estimated cost |
-| 🗑️ | **Clean terminate** | Stops instances, removes SSH aliases, shows shutting-down state until fully gone |
+| 💾 | **EBS data volumes** | Attach persistent storage at `/data` — survives spot interruptions and termination, reattach to new instances |
+| 🗑️ | **Clean terminate** | Stops instances, removes SSH aliases, cleans up EBS volumes (or preserves with `--keep-ebs`) |
 
 ### 🎯 Target Workflows
 
@@ -113,16 +114,24 @@ aws-bootstrap launch --python-version 3.13
 # Use a non-default SSH port
 aws-bootstrap launch --ssh-port 2222
 
+# Attach a persistent EBS data volume (96 GB gp3, mounted at /data)
+aws-bootstrap launch --ebs-storage 96
+
+# Reattach an existing EBS volume from a previous instance
+aws-bootstrap launch --ebs-volume-id vol-0abc123def456
+
 # Use a specific AWS profile
 aws-bootstrap launch --profile my-aws-profile
 ```
 
 After launch, the CLI:
 
-1. **Adds an SSH alias** (e.g. `aws-gpu1`) to `~/.ssh/config`
-2. **Runs remote setup** — installs utilities, creates a Python venv, installs CUDA-matched PyTorch, sets up Jupyter
-3. **Runs a CUDA smoke test** — verifies `torch.cuda.is_available()` and runs a quick GPU matmul
-4. **Prints connection commands** — SSH, Jupyter tunnel, GPU benchmark, and terminate
+1. **Creates/attaches EBS volume** (if `--ebs-storage` or `--ebs-volume-id` was specified)
+2. **Adds an SSH alias** (e.g. `aws-gpu1`) to `~/.ssh/config`
+3. **Runs remote setup** — installs utilities, creates a Python venv, installs CUDA-matched PyTorch, sets up Jupyter
+4. **Mounts EBS volume** at `/data` (if applicable — formats new volumes, mounts existing ones as-is)
+5. **Runs a CUDA smoke test** — verifies `torch.cuda.is_available()` and runs a quick GPU matmul
+6. **Prints connection commands** — SSH, Jupyter tunnel, GPU benchmark, and terminate
 
 ```bash
 ssh aws-gpu1                  # venv auto-activates on login
@@ -135,7 +144,7 @@ The setup script runs automatically on the instance after SSH becomes available:
 | Step | What |
 |------|------|
 | **GPU verify** | Confirms `nvidia-smi` and `nvcc` are working |
-| **Utilities** | Installs `htop`, `tmux`, `tree`, `jq` |
+| **Utilities** | Installs `htop`, `tmux`, `tree`, `jq`, `ffmpeg` |
 | **Python venv** | Creates `~/venv` with `uv`, auto-activates in `~/.bashrc`. Use `--python-version` to pin a specific Python (e.g. `3.13`) |
 | **CUDA-aware PyTorch** | Detects CUDA toolkit version → installs PyTorch from the matching `cu{TAG}` wheel index |
 | **CUDA smoke test** | Runs `torch.cuda.is_available()` + GPU matmul to verify the stack |
@@ -242,6 +251,9 @@ aws-bootstrap status --region us-east-1
 # Terminate all aws-bootstrap instances (with confirmation prompt)
 aws-bootstrap terminate
 
+# Terminate but preserve EBS data volumes for reuse
+aws-bootstrap terminate --keep-ebs
+
 # Terminate by SSH alias (resolved via ~/.ssh/config)
 aws-bootstrap terminate aws-gpu1
 
@@ -253,6 +265,15 @@ aws-bootstrap terminate aws-gpu1 i-def456
 
 # Skip confirmation prompt
 aws-bootstrap terminate --yes
+
+# Remove stale SSH config entries for terminated instances
+aws-bootstrap cleanup
+
+# Preview what would be removed without modifying config
+aws-bootstrap cleanup --dry-run
+
+# Skip confirmation prompt
+aws-bootstrap cleanup --yes
 ```
 
 `status --gpu` reports both the **installed CUDA toolkit** version (from `nvcc`) and the **maximum CUDA version supported by the driver** (from `nvidia-smi`), so you can see at a glance whether they match:
@@ -263,6 +284,31 @@ CUDA: 12.8 (driver supports up to 13.0)
 
 SSH aliases are managed automatically — they're created on `launch`, shown in `status`, and cleaned up on `terminate`. Aliases use sequential numbering (`aws-gpu1`, `aws-gpu2`, etc.) and never reuse numbers from previous instances. You can use aliases anywhere you'd use an instance ID, e.g. `aws-bootstrap terminate aws-gpu1`.
 
+## EBS Data Volumes
+
+Attach persistent EBS storage to keep datasets and model checkpoints across instance lifecycles. Volumes are mounted at `/data` and persist independently of the instance.
+
+```bash
+# Create a new 96 GB gp3 volume, formatted and mounted at /data
+aws-bootstrap launch --ebs-storage 96
+
+# After terminating with --keep-ebs, reattach the same volume to a new instance
+aws-bootstrap terminate --keep-ebs
+# Output: Preserving EBS volume: vol-0abc123...
+#         Reattach with: aws-bootstrap launch --ebs-volume-id vol-0abc123...
+
+aws-bootstrap launch --ebs-volume-id vol-0abc123def456
+```
+
+Key behaviors:
+- `--ebs-storage` and `--ebs-volume-id` are mutually exclusive
+- New volumes are formatted as ext4; existing volumes are mounted as-is
+- Volumes are tagged for automatic discovery by `status` and `terminate`
+- `terminate` deletes data volumes by default; use `--keep-ebs` to preserve them
+- **Spot-safe** — data volumes survive spot interruptions. If AWS reclaims your instance, the volume detaches automatically and can be reattached to a new instance with `--ebs-volume-id`
+- EBS volumes must be in the same availability zone as the instance
+- Mount failures are non-fatal — the instance remains usable
+
 ## EC2 vCPU Quotas
 
 AWS accounts have [service quotas](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-resource-limits.html) that limit how many vCPUs you can run per instance family. New or lightly-used accounts often have a **default quota of 0 vCPUs** for GPU instance families (G and VT), which will cause errors on launch:

{aws_bootstrap_g4dn-0.5.0 → aws_bootstrap_g4dn-0.6.0}/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,17 @@ from .ec2 import (
     list_amis,
     list_instance_types,
     terminate_tagged_instances,
+    validate_ebs_volume,
     wait_instance_ready,
 )
 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,
@@ -120,6 +129,18 @@ 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.",
+)
 def launch(
     instance_type,
     ami_filter,
@@ -134,8 +155,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 +174,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,18 +190,21 @@ 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"
@@ -193,18 +224,22 @@ def launch(
             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")
         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 +248,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
 
-    # Step 6: SSH and remote setup
-    step(6, 6, "Waiting for SSH access...")
+        info(f"Attaching {ebs_volume_attached} to {instance_id}...")
+        attach_ebs_volume(ec2, ebs_volume_attached, instance_id)
+        success("EBS volume attached.")
+
+    # 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 +303,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,
@@ -260,6 +341,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 ""
 
@@ -371,6 +458,12 @@ def status(region, profile, gpu, instructions):
             else:
                 click.echo("    GPU: " + click.style("unavailable", dim=True))
 
+        # EBS data volumes
+        ebs_volumes = find_ebs_volumes_for_instance(ec2, inst["InstanceId"], "aws-bootstrap-g4dn")
+        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})")
+
         lifecycle = inst["Lifecycle"]
         is_spot = lifecycle == "spot"
 
@@ -429,8 +522,9 @@ 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):
+def terminate(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,
@@ -468,6 +562,13 @@ 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()
     for change in changes:
@@ -479,10 +580,73 @@ def terminate(region, profile, yes, instance_ids):
         removed_alias = remove_ssh_host(change["InstanceId"])
         if removed_alias:
             info(f"Removed SSH config alias: {removed_alias}")
+
+    # Handle EBS volume cleanup
+    for _iid, volumes in ebs_by_instance.items():
+        for vol in volumes:
+            vid = vol["VolumeId"]
+            if keep_ebs:
+                click.echo()
+                info(f"Preserving EBS volume: {vid} ({vol['Size']} GB)")
+                info(f"Reattach with: aws-bootstrap launch --ebs-volume-id {vid}")
+            else:
+                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}")
+                except Exception as e:
+                    warn(f"Failed to delete EBS volume {vid}: {e}")
+
     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.")
+def cleanup(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")
+
+    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:
+        click.secho("No stale SSH config entries found.", fg="green")
+        return
+
+    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:
+        click.echo()
+        for iid, alias in stale:
+            info(f"Would remove {alias} ({iid})")
+        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)
+    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
 # ---------------------------------------------------------------------------

{aws_bootstrap_g4dn-0.5.0 → aws_bootstrap_g4dn-0.6.0}/aws_bootstrap/config.py

@@ -24,3 +24,5 @@ class LaunchConfig:
     alias_prefix: str = "aws-gpu"
     ssh_port: int = 22
     python_version: str | None = None
+    ebs_storage: int | None = None
+    ebs_volume_id: str | None = None