aws-bootstrap-g4dn 0.5.0__py3-none-any.whl → 0.6.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,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/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

aws_bootstrap/ec2.py

@@ -9,6 +9,10 @@ import click
 from .config import LaunchConfig
 
 
+EBS_DEVICE_NAME = "/dev/sdf"
+EBS_MOUNT_POINT = "/data"
+
+
 class CLIError(click.ClickException):
     """A ClickException that displays the error message in red."""
 
@@ -339,3 +343,127 @@ def wait_instance_ready(ec2_client, instance_id: str) -> dict:
     desc = ec2_client.describe_instances(InstanceIds=[instance_id])
     instance = desc["Reservations"][0]["Instances"][0]
     return instance
+
+
+# ---------------------------------------------------------------------------
+# EBS data volume operations
+# ---------------------------------------------------------------------------
+
+
+def create_ebs_volume(ec2_client, size_gb: int, availability_zone: str, tag_value: str, instance_id: str) -> str:
+    """Create a gp3 EBS volume and wait for it to become available.
+
+    Returns the volume ID.
+    """
+    response = ec2_client.create_volume(
+        AvailabilityZone=availability_zone,
+        Size=size_gb,
+        VolumeType="gp3",
+        TagSpecifications=[
+            {
+                "ResourceType": "volume",
+                "Tags": [
+                    {"Key": "created-by", "Value": tag_value},
+                    {"Key": "Name", "Value": f"aws-bootstrap-data-{instance_id}"},
+                    {"Key": "aws-bootstrap-instance", "Value": instance_id},
+                ],
+            }
+        ],
+    )
+    volume_id = response["VolumeId"]
+
+    waiter = ec2_client.get_waiter("volume_available")
+    waiter.wait(VolumeIds=[volume_id], WaiterConfig={"Delay": 5, "MaxAttempts": 24})
+    return volume_id
+
+
+def validate_ebs_volume(ec2_client, volume_id: str, availability_zone: str) -> dict:
+    """Validate that an existing EBS volume can be attached.
+
+    Checks that the volume exists, is available (not in-use), and is in the
+    correct availability zone. Returns the volume description dict.
+
+    Raises CLIError for validation failures.
+    """
+    try:
+        response = ec2_client.describe_volumes(VolumeIds=[volume_id])
+    except botocore.exceptions.ClientError as e:
+        if e.response["Error"]["Code"] == "InvalidVolume.NotFound":
+            raise CLIError(f"EBS volume not found: {volume_id}") from None
+        raise
+
+    volumes = response["Volumes"]
+    if not volumes:
+        raise CLIError(f"EBS volume not found: {volume_id}")
+
+    vol = volumes[0]
+
+    if vol["State"] != "available":
+        raise CLIError(
+            f"EBS volume {volume_id} is currently '{vol['State']}' (must be 'available').\n"
+            "  Detach it from its current instance first."
+        )
+
+    if vol["AvailabilityZone"] != availability_zone:
+        raise CLIError(
+            f"EBS volume {volume_id} is in {vol['AvailabilityZone']} "
+            f"but the instance is in {availability_zone}.\n"
+            "  EBS volumes must be in the same availability zone as the instance."
+        )
+
+    return vol
+
+
+def attach_ebs_volume(ec2_client, volume_id: str, instance_id: str, device_name: str = EBS_DEVICE_NAME) -> None:
+    """Attach an EBS volume to an instance and wait for it to be in-use."""
+    ec2_client.attach_volume(
+        VolumeId=volume_id,
+        InstanceId=instance_id,
+        Device=device_name,
+    )
+    waiter = ec2_client.get_waiter("volume_in_use")
+    waiter.wait(VolumeIds=[volume_id], WaiterConfig={"Delay": 5, "MaxAttempts": 24})
+
+
+def detach_ebs_volume(ec2_client, volume_id: str) -> None:
+    """Detach an EBS volume and wait for it to become available."""
+    ec2_client.detach_volume(VolumeId=volume_id)
+    waiter = ec2_client.get_waiter("volume_available")
+    waiter.wait(VolumeIds=[volume_id], WaiterConfig={"Delay": 5, "MaxAttempts": 24})
+
+
+def delete_ebs_volume(ec2_client, volume_id: str) -> None:
+    """Delete an EBS volume."""
+    ec2_client.delete_volume(VolumeId=volume_id)
+
+
+def find_ebs_volumes_for_instance(ec2_client, instance_id: str, tag_value: str) -> list[dict]:
+    """Find EBS data volumes associated with an instance via tags.
+
+    Returns a list of dicts with VolumeId, Size, Device, and State.
+    Excludes root volumes (only returns volumes tagged by aws-bootstrap).
+    """
+    try:
+        response = ec2_client.describe_volumes(
+            Filters=[
+                {"Name": "tag:aws-bootstrap-instance", "Values": [instance_id]},
+                {"Name": "tag:created-by", "Values": [tag_value]},
+            ]
+        )
+    except botocore.exceptions.ClientError:
+        return []
+
+    volumes = []
+    for vol in response.get("Volumes", []):
+        device = ""
+        if vol.get("Attachments"):
+            device = vol["Attachments"][0].get("Device", "")
+        volumes.append(
+            {
+                "VolumeId": vol["VolumeId"],
+                "Size": vol["Size"],
+                "Device": device,
+                "State": vol["State"],
+            }
+        )
+    return volumes

aws_bootstrap/resources/remote_setup.sh

@@ -48,8 +48,8 @@ fi
 # 2. Install utilities
 echo ""
 echo "[2/6] Installing utilities..."
-sudo apt-get update -qq
-sudo apt-get install -y -qq htop tmux tree jq
+sudo DEBIAN_FRONTEND=noninteractive apt-get update -qq
+sudo DEBIAN_FRONTEND=noninteractive apt-get install -y -qq htop tmux tree jq ffmpeg
 
 # 3. Set up Python environment with uv
 echo ""

aws_bootstrap/ssh.py

@@ -374,6 +374,37 @@ def list_ssh_hosts(config_path: Path | None = None) -> dict[str, str]:
     return result
 
 
+def find_stale_ssh_hosts(live_instance_ids: set[str], config_path: Path | None = None) -> list[tuple[str, str]]:
+    """Identify SSH config entries whose instances no longer exist.
+
+    Returns ``[(instance_id, alias), ...]`` for entries where the instance ID
+    is **not** in *live_instance_ids*, sorted by alias.
+    """
+    hosts = list_ssh_hosts(config_path)
+    stale = [(iid, alias) for iid, alias in hosts.items() if iid not in live_instance_ids]
+    stale.sort(key=lambda t: t[1])
+    return stale
+
+
+def cleanup_stale_ssh_hosts(
+    live_instance_ids: set[str],
+    config_path: Path | None = None,
+    dry_run: bool = False,
+) -> list[CleanupResult]:
+    """Remove SSH config entries for terminated/non-existent instances.
+
+    If *dry_run* is ``True``, entries are identified but not removed.
+    Returns a list of :class:`CleanupResult` objects.
+    """
+    stale = find_stale_ssh_hosts(live_instance_ids, config_path)
+    results: list[CleanupResult] = []
+    for iid, alias in stale:
+        if not dry_run:
+            remove_ssh_host(iid, config_path)
+        results.append(CleanupResult(instance_id=iid, alias=alias, removed=not dry_run))
+    return results
+
+
 _INSTANCE_ID_RE = re.compile(r"^i-[0-9a-f]{8,17}$")
 
 
@@ -402,6 +433,15 @@ def resolve_instance_id(value: str, config_path: Path | None = None) -> str | No
     return None
 
 
+@dataclass
+class CleanupResult:
+    """Result of cleaning up a single stale SSH config entry."""
+
+    instance_id: str
+    alias: str
+    removed: bool
+
+
 @dataclass
 class SSHHostDetails:
     """Connection details parsed from an SSH config stanza."""
@@ -515,6 +555,87 @@ def query_gpu_info(host: str, user: str, key_path: Path, timeout: int = 10, port
         return None
 
 
+# ---------------------------------------------------------------------------
+# EBS volume mount
+# ---------------------------------------------------------------------------
+
+
+def mount_ebs_volume(
+    host: str,
+    user: str,
+    key_path: Path,
+    volume_id: str,
+    mount_point: str = "/data",
+    format_volume: bool = True,
+    port: int = 22,
+) -> bool:
+    """Mount an EBS volume on the remote instance via SSH.
+
+    Detects the NVMe device by volume ID serial, formats if requested,
+    mounts at *mount_point*, and adds an fstab entry for persistence.
+
+    Returns True on success, False on failure.
+    """
+    ssh_opts = _ssh_opts(key_path)
+    port_opts = ["-p", str(port)] if port != 22 else []
+
+    # Strip the vol- prefix and hyphen for NVMe serial matching
+    vol_serial = volume_id.replace("-", "")
+
+    format_cmd = ""
+    if format_volume:
+        format_cmd = (
+            '  if ! sudo blkid "$DEVICE" > /dev/null 2>&1; then\n'
+            '    echo "Formatting $DEVICE as ext4..."\n'
+            '    sudo mkfs.ext4 "$DEVICE"\n'
+            "  fi\n"
+        )
+
+    remote_script = (
+        "set -e\n"
+        "# Detect EBS device by NVMe serial (Nitro instances)\n"
+        f'SERIAL="{vol_serial}"\n'
+        "DEVICE=$(lsblk -o NAME,SERIAL -dpn 2>/dev/null | "
+        "awk -v s=\"$SERIAL\" '$2 == s {print $1}' | head -1)\n"
+        "# Fallback to common device paths\n"
+        'if [ -z "$DEVICE" ]; then\n'
+        "  for dev in /dev/nvme1n1 /dev/xvdf /dev/sdf; do\n"
+        '    if [ -b "$dev" ]; then DEVICE="$dev"; break; fi\n'
+        "  done\n"
+        "fi\n"
+        'if [ -z "$DEVICE" ]; then\n'
+        '  echo "ERROR: Could not find EBS device" >&2\n'
+        "  exit 1\n"
+        "fi\n"
+        'echo "Found EBS device: $DEVICE"\n'
+        f"{format_cmd}"
+        f"sudo mkdir -p {mount_point}\n"
+        f'sudo mount "$DEVICE" {mount_point}\n'
+        f"sudo chown {user}:{user} {mount_point}\n"
+        "# Add fstab entry for reboot persistence\n"
+        'UUID=$(sudo blkid -s UUID -o value "$DEVICE")\n'
+        'if [ -n "$UUID" ]; then\n'
+        f'  if ! grep -q "$UUID" /etc/fstab; then\n'
+        f'    echo "UUID=$UUID {mount_point} ext4 defaults,nofail 0 2" | sudo tee -a /etc/fstab > /dev/null\n'
+        "  fi\n"
+        "fi\n"
+        f'echo "Mounted $DEVICE at {mount_point}"'
+    )
+
+    cmd = [
+        "ssh",
+        *ssh_opts,
+        *port_opts,
+        "-o",
+        "ConnectTimeout=10",
+        f"{user}@{host}",
+        remote_script,
+    ]
+
+    result = subprocess.run(cmd, capture_output=False)
+    return result.returncode == 0
+
+
 # ---------------------------------------------------------------------------
 # Internal helpers
 # ---------------------------------------------------------------------------