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

aws_bootstrap/tests/test_ssh_config.py

@@ -10,7 +10,9 @@ from aws_bootstrap.ssh import (
     _next_alias,
     _read_ssh_config,
     add_ssh_host,
+    cleanup_stale_ssh_hosts,
     find_ssh_alias,
+    find_stale_ssh_hosts,
     get_ssh_host_details,
     list_ssh_hosts,
     remove_ssh_host,
@@ -407,3 +409,77 @@ def test_resolve_unknown_alias_returns_none(tmp_path):
 def test_resolve_nonexistent_config_returns_none(tmp_path):
     cfg = tmp_path / "no_such_file"
     assert resolve_instance_id("aws-gpu1", config_path=cfg) is None
+
+
+# ---------------------------------------------------------------------------
+# Cleanup: find_stale_ssh_hosts / cleanup_stale_ssh_hosts
+# ---------------------------------------------------------------------------
+
+
+def test_find_stale_ssh_hosts_finds_orphans(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-111aaaa1", "1.1.1.1", "ubuntu", KEY_PATH, config_path=cfg)
+    add_ssh_host("i-222bbbb2", "2.2.2.2", "ubuntu", KEY_PATH, config_path=cfg)
+    stale = find_stale_ssh_hosts({"i-111aaaa1"}, config_path=cfg)
+    assert stale == [("i-222bbbb2", "aws-gpu2")]
+
+
+def test_find_stale_ssh_hosts_none_stale(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-111aaaa1", "1.1.1.1", "ubuntu", KEY_PATH, config_path=cfg)
+    add_ssh_host("i-222bbbb2", "2.2.2.2", "ubuntu", KEY_PATH, config_path=cfg)
+    stale = find_stale_ssh_hosts({"i-111aaaa1", "i-222bbbb2"}, config_path=cfg)
+    assert stale == []
+
+
+def test_find_stale_ssh_hosts_all_stale(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-111aaaa1", "1.1.1.1", "ubuntu", KEY_PATH, config_path=cfg)
+    add_ssh_host("i-222bbbb2", "2.2.2.2", "ubuntu", KEY_PATH, config_path=cfg)
+    stale = find_stale_ssh_hosts(set(), config_path=cfg)
+    assert len(stale) == 2
+    assert ("i-111aaaa1", "aws-gpu1") in stale
+    assert ("i-222bbbb2", "aws-gpu2") in stale
+
+
+def test_find_stale_ssh_hosts_empty_config(tmp_path):
+    cfg = _config_path(tmp_path)
+    cfg.parent.mkdir(parents=True, exist_ok=True)
+    cfg.write_text("")
+    stale = find_stale_ssh_hosts(set(), config_path=cfg)
+    assert stale == []
+
+
+def test_cleanup_stale_ssh_hosts_removes(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-111aaaa1", "1.1.1.1", "ubuntu", KEY_PATH, config_path=cfg)
+    add_ssh_host("i-222bbbb2", "2.2.2.2", "ubuntu", KEY_PATH, config_path=cfg)
+    results = cleanup_stale_ssh_hosts({"i-111aaaa1"}, config_path=cfg)
+    assert len(results) == 1
+    assert results[0].instance_id == "i-222bbbb2"
+    assert results[0].alias == "aws-gpu2"
+    assert results[0].removed is True
+    # Verify it was actually removed from the config
+    content = cfg.read_text()
+    assert "i-222bbbb2" not in content
+    assert "i-111aaaa1" in content
+
+
+def test_cleanup_stale_ssh_hosts_dry_run(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-111aaaa1", "1.1.1.1", "ubuntu", KEY_PATH, config_path=cfg)
+    add_ssh_host("i-222bbbb2", "2.2.2.2", "ubuntu", KEY_PATH, config_path=cfg)
+    results = cleanup_stale_ssh_hosts({"i-111aaaa1"}, config_path=cfg, dry_run=True)
+    assert len(results) == 1
+    assert results[0].removed is False
+    # Verify config is unchanged
+    content = cfg.read_text()
+    assert "i-222bbbb2" in content
+    assert "i-111aaaa1" in content
+
+
+def test_cleanup_stale_ssh_hosts_no_stale(tmp_path):
+    cfg = _config_path(tmp_path)
+    add_ssh_host("i-111aaaa1", "1.1.1.1", "ubuntu", KEY_PATH, config_path=cfg)
+    results = cleanup_stale_ssh_hosts({"i-111aaaa1"}, config_path=cfg)
+    assert results == []

aws_bootstrap/tests/test_ssh_ebs.py

@@ -0,0 +1,76 @@
+"""Tests for mount_ebs_volume SSH function."""
+
+from __future__ import annotations
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+from aws_bootstrap.ssh import mount_ebs_volume
+
+
+KEY_PATH = Path("/home/user/.ssh/id_ed25519.pub")
+
+
+@patch("aws_bootstrap.ssh.subprocess.run")
+def test_mount_ebs_volume_success_format(mock_run):
+    """New volume: SSH command includes mkfs."""
+    mock_run.return_value = MagicMock(returncode=0)
+
+    result = mount_ebs_volume("1.2.3.4", "ubuntu", KEY_PATH, "vol-abc123", format_volume=True)
+
+    assert result is True
+    mock_run.assert_called_once()
+    cmd = mock_run.call_args[0][0]
+    script = cmd[-1]
+    assert "mkfs.ext4" in script
+    assert "/data" in script
+    assert "volabc123" in script  # stripped vol- hyphen
+
+
+@patch("aws_bootstrap.ssh.subprocess.run")
+def test_mount_ebs_volume_success_no_format(mock_run):
+    """Existing volume: SSH command skips mkfs."""
+    mock_run.return_value = MagicMock(returncode=0)
+
+    result = mount_ebs_volume("1.2.3.4", "ubuntu", KEY_PATH, "vol-abc123", format_volume=False)
+
+    assert result is True
+    mock_run.assert_called_once()
+    cmd = mock_run.call_args[0][0]
+    script = cmd[-1]
+    assert "mkfs" not in script
+    assert "/data" in script
+
+
+@patch("aws_bootstrap.ssh.subprocess.run")
+def test_mount_ebs_volume_failure(mock_run):
+    """Non-zero exit code returns False."""
+    mock_run.return_value = MagicMock(returncode=1)
+
+    result = mount_ebs_volume("1.2.3.4", "ubuntu", KEY_PATH, "vol-abc123")
+
+    assert result is False
+
+
+@patch("aws_bootstrap.ssh.subprocess.run")
+def test_mount_ebs_volume_custom_port(mock_run):
+    """Non-default port is passed as -p flag."""
+    mock_run.return_value = MagicMock(returncode=0)
+
+    mount_ebs_volume("1.2.3.4", "ubuntu", KEY_PATH, "vol-abc123", port=2222)
+
+    cmd = mock_run.call_args[0][0]
+    assert "-p" in cmd
+    port_idx = cmd.index("-p")
+    assert cmd[port_idx + 1] == "2222"
+
+
+@patch("aws_bootstrap.ssh.subprocess.run")
+def test_mount_ebs_volume_custom_mount_point(mock_run):
+    """Custom mount point appears in the SSH script."""
+    mock_run.return_value = MagicMock(returncode=0)
+
+    mount_ebs_volume("1.2.3.4", "ubuntu", KEY_PATH, "vol-abc123", mount_point="/mnt/data")
+
+    cmd = mock_run.call_args[0][0]
+    script = cmd[-1]
+    assert "/mnt/data" in script

{aws_bootstrap_g4dn-0.5.0.dist-info → aws_bootstrap_g4dn-0.6.0.dist-info}/METADATA

@@ -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.dist-info → aws_bootstrap_g4dn-0.6.0.dist-info}/RECORD

@@ -1,27 +1,29 @@
 aws_bootstrap/__init__.py,sha256=kl_jvrunGyIyizdRqAP6ROb5P1BBrXX5PTq5gq1ipU0,82
-aws_bootstrap/cli.py,sha256=N2hT0XEC-4k5Cs3iGfA_xt_onc__NMNmh8fCaV4frgc,21076
-aws_bootstrap/config.py,sha256=TeCOYDlijT-KD5SFIzc-VvBhOqcq9YCgen9NK63rka8,895
-aws_bootstrap/ec2.py,sha256=LHpzW91ayK45gsWV_B4LanSZIhWggqTsL31qHUceiaA,12274
+aws_bootstrap/cli.py,sha256=n3Ep_7zhBiRSU4ZUeGVqTRb81nzo98mxzQSKdAuiopY,27788
+aws_bootstrap/config.py,sha256=p770XgjfuK1-wVkAEeBdtJSVkc58DKFHgaJlZ-zbGmk,967
+aws_bootstrap/ec2.py,sha256=uNqxWWfPfGCbujQ3eonvqjjxLE76fsEyNchPS6byR6c,16719
 aws_bootstrap/gpu.py,sha256=WTnHR0s3mQHDlnzqRgqAC6omWz7nT5YtGpcs0Bf88jk,692
-aws_bootstrap/ssh.py,sha256=0acHNX7IG6PUvp6T72l9kHTwUs5sVXFAyJXvUfA3qnE,20131
+aws_bootstrap/ssh.py,sha256=xY0Yn5q4aA0Xb3ejNY-KCbooZArXRGpimSnbJiBLI_w,24059
 aws_bootstrap/resources/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aws_bootstrap/resources/gpu_benchmark.py,sha256=1eFt_3MXvoLhs9HahrRPhbxvtdjFaXG2Ty3GEg7Gud0,29366
 aws_bootstrap/resources/gpu_smoke_test.ipynb,sha256=XvAOEIPa5H9ri5mRZqOdknmwOwKNvCME6DzBGuhRYfg,10698
 aws_bootstrap/resources/launch.json,sha256=ZOcvHLy3-zBOqRTtFzuyn-_2tB64yuEn8PrJOoZ-PgE,1484
-aws_bootstrap/resources/remote_setup.sh,sha256=z_YGdzwEHWInkE3dZVbBNa0F_joTeVhnOpCYOj1CK30,8331
+aws_bootstrap/resources/remote_setup.sh,sha256=i9qXzAqYrnDWt6nFrqeghYWVHmbbzP-OS7O_YJB3GYU,8400
 aws_bootstrap/resources/requirements.txt,sha256=gpYl1MFCfWXiAhbIUgAjuTHONz3MKci25msIyOkMmUk,75
 aws_bootstrap/resources/saxpy.cu,sha256=1BSESEwGGCx3KWx9ZJ8jiPHQ42KzQN6i2aP0I28bPsA,1178
 aws_bootstrap/resources/tasks.json,sha256=6U8pB1N8YIWgUCfFet4ne3nYnI92tWv5D5kPiQG3Zlg,1576
 aws_bootstrap/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aws_bootstrap/tests/test_cli.py,sha256=m_4tIX0xYZ8BxDkHPGAWiPAKo4vETaTSKSJbyn3K1Cg,34731
-aws_bootstrap/tests/test_config.py,sha256=arvET6KNl4Vqsz0zFrSdhciXGU688bfsvCr3dSpziN0,1050
+aws_bootstrap/tests/test_cli.py,sha256=OOYzopo8gkp6fO1MT10fsqAhIs5pN9cCpJyOzWNXHCg,48638
+aws_bootstrap/tests/test_config.py,sha256=vspSGoben_i7m4Fh6UGSes6Fkr789Y1eaOLe54fRSGc,1524
+aws_bootstrap/tests/test_ebs.py,sha256=B2HrgSmS7yroz6zzRuPxKIXmQGlWesuGqOtybyZmHJQ,7582
 aws_bootstrap/tests/test_ec2.py,sha256=Jmqsjv973hxXbZWfGgECtm6aa2156Lzji227sYMBuMg,10547
 aws_bootstrap/tests/test_gpu.py,sha256=rbMuda_sIVbaCzkWXoLv9YIfnWztgRoP7NuVL8XHrUY,3871
-aws_bootstrap/tests/test_ssh_config.py,sha256=YYtv82zBBLGioTo58iC31_5jUli1s0eoGV9VRCobOgY,14059
+aws_bootstrap/tests/test_ssh_config.py,sha256=qy3UDdvkTfrALiF-W3m8aKvnQj3BeCrZdLjG75tcVJU,17131
+aws_bootstrap/tests/test_ssh_ebs.py,sha256=ipt0xOzdf3kfkVt42Dgr_z7D6JDIMuRi3DqX0OP8sm0,2342
 aws_bootstrap/tests/test_ssh_gpu.py,sha256=dRp86Og-8GqiATSff3rxhu83mBZdGgqI4UOnoC00Ln0,1454
-aws_bootstrap_g4dn-0.5.0.dist-info/licenses/LICENSE,sha256=Hen77Mt8sazSQJ9DgrmZuAvDwo2vc5JAkR_avuFV-CM,1067
-aws_bootstrap_g4dn-0.5.0.dist-info/METADATA,sha256=t8m53ZodJlZyMffeSu3Wk5bMt-Dm_Jl3q_HTbRLQbYE,13728
-aws_bootstrap_g4dn-0.5.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-aws_bootstrap_g4dn-0.5.0.dist-info/entry_points.txt,sha256=T8FXfOgmLEvFi8DHaFJ3tCzId9J3_d2Y6qT98OXxCjA,57
-aws_bootstrap_g4dn-0.5.0.dist-info/top_level.txt,sha256=mix9gZRs8JUv0OMSB_rwdGcRnTKzsKgHrE5fyAn5zJw,14
-aws_bootstrap_g4dn-0.5.0.dist-info/RECORD,,
+aws_bootstrap_g4dn-0.6.0.dist-info/licenses/LICENSE,sha256=Hen77Mt8sazSQJ9DgrmZuAvDwo2vc5JAkR_avuFV-CM,1067
+aws_bootstrap_g4dn-0.6.0.dist-info/METADATA,sha256=Ot9yCJfJup1ZzW-0cq99zT9bbswHnRJ4SWxEsJ-pK58,15859
+aws_bootstrap_g4dn-0.6.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+aws_bootstrap_g4dn-0.6.0.dist-info/entry_points.txt,sha256=T8FXfOgmLEvFi8DHaFJ3tCzId9J3_d2Y6qT98OXxCjA,57
+aws_bootstrap_g4dn-0.6.0.dist-info/top_level.txt,sha256=mix9gZRs8JUv0OMSB_rwdGcRnTKzsKgHrE5fyAn5zJw,14
+aws_bootstrap_g4dn-0.6.0.dist-info/RECORD,,