aws-bootstrap-g4dn 0.2.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

aws_bootstrap/ssh.py

@@ -12,6 +12,8 @@ from pathlib import Path
 
 import click
 
+from .gpu import _GPU_ARCHITECTURES, GpuInfo
+
 
 # ---------------------------------------------------------------------------
 # SSH config markers
@@ -72,17 +74,18 @@ def import_key_pair(ec2_client, key_name: str, key_path: Path) -> str:
     return key_name
 
 
-def wait_for_ssh(host: str, user: str, key_path: Path, retries: int = 30, delay: int = 10) -> bool:
+def wait_for_ssh(host: str, user: str, key_path: Path, retries: int = 30, delay: int = 10, port: int = 22) -> bool:
     """Wait for SSH to become available on the instance.
 
-    Tries a TCP connection to port 22 first, then an actual SSH command.
+    Tries a TCP connection to the SSH port first, then an actual SSH command.
     """
     base_opts = _ssh_opts(key_path)
+    port_opts = ["-p", str(port)] if port != 22 else []
 
     for attempt in range(1, retries + 1):
-        # First check if port 22 is open
+        # First check if the SSH port is open
         try:
-            sock = socket.create_connection((host, 22), timeout=5)
+            sock = socket.create_connection((host, port), timeout=5)
             sock.close()
         except (TimeoutError, ConnectionRefusedError, OSError):
             click.echo("  SSH not ready " + click.style(f"(attempt {attempt}/{retries})", dim=True) + ", waiting...")
@@ -90,11 +93,18 @@ def wait_for_ssh(host: str, user: str, key_path: Path, retries: int = 30, delay:
             continue
 
         # Port is open, try actual SSH
-        result = subprocess.run(
-            ["ssh", *base_opts, "-o", "ConnectTimeout=10", "-o", "BatchMode=yes", f"{user}@{host}", "echo ok"],
-            capture_output=True,
-            text=True,
-        )
+        cmd = [
+            "ssh",
+            *base_opts,
+            *port_opts,
+            "-o",
+            "ConnectTimeout=10",
+            "-o",
+            "BatchMode=yes",
+            f"{user}@{host}",
+            "echo ok",
+        ]
+        result = subprocess.run(cmd, capture_output=True, text=True)
         if result.returncode == 0:
             click.secho("  SSH connection established.", fg="green")
             return True
@@ -105,15 +115,19 @@ def wait_for_ssh(host: str, user: str, key_path: Path, retries: int = 30, delay:
     return False
 
 
-def run_remote_setup(host: str, user: str, key_path: Path, script_path: Path) -> bool:
+def run_remote_setup(
+    host: str, user: str, key_path: Path, script_path: Path, python_version: str | None = None, port: int = 22
+) -> bool:
     """SCP the setup script and requirements.txt to the instance and execute."""
     ssh_opts = _ssh_opts(key_path)
+    scp_port_opts = ["-P", str(port)] if port != 22 else []
+    ssh_port_opts = ["-p", str(port)] if port != 22 else []
     requirements_path = script_path.parent / "requirements.txt"
 
     # SCP the requirements file
     click.echo("  Uploading requirements.txt...")
     req_result = subprocess.run(
-        ["scp", *ssh_opts, str(requirements_path), f"{user}@{host}:/tmp/requirements.txt"],
+        ["scp", *ssh_opts, *scp_port_opts, str(requirements_path), f"{user}@{host}:/tmp/requirements.txt"],
         capture_output=True,
         text=True,
     )
@@ -125,7 +139,7 @@ def run_remote_setup(host: str, user: str, key_path: Path, script_path: Path) ->
     benchmark_path = script_path.parent / "gpu_benchmark.py"
     click.echo("  Uploading gpu_benchmark.py...")
     bench_result = subprocess.run(
-        ["scp", *ssh_opts, str(benchmark_path), f"{user}@{host}:/tmp/gpu_benchmark.py"],
+        ["scp", *ssh_opts, *scp_port_opts, str(benchmark_path), f"{user}@{host}:/tmp/gpu_benchmark.py"],
         capture_output=True,
         text=True,
     )
@@ -137,7 +151,7 @@ def run_remote_setup(host: str, user: str, key_path: Path, script_path: Path) ->
     notebook_path = script_path.parent / "gpu_smoke_test.ipynb"
     click.echo("  Uploading gpu_smoke_test.ipynb...")
     nb_result = subprocess.run(
-        ["scp", *ssh_opts, str(notebook_path), f"{user}@{host}:/tmp/gpu_smoke_test.ipynb"],
+        ["scp", *ssh_opts, *scp_port_opts, str(notebook_path), f"{user}@{host}:/tmp/gpu_smoke_test.ipynb"],
         capture_output=True,
         text=True,
     )
@@ -145,10 +159,46 @@ def run_remote_setup(host: str, user: str, key_path: Path, script_path: Path) ->
         click.secho(f"  SCP failed: {nb_result.stderr}", fg="red", err=True)
         return False
 
+    # SCP the CUDA example source
+    saxpy_path = script_path.parent / "saxpy.cu"
+    click.echo("  Uploading saxpy.cu...")
+    saxpy_result = subprocess.run(
+        ["scp", *ssh_opts, *scp_port_opts, str(saxpy_path), f"{user}@{host}:/tmp/saxpy.cu"],
+        capture_output=True,
+        text=True,
+    )
+    if saxpy_result.returncode != 0:
+        click.secho(f"  SCP failed: {saxpy_result.stderr}", fg="red", err=True)
+        return False
+
+    # SCP the VSCode launch.json
+    launch_json_path = script_path.parent / "launch.json"
+    click.echo("  Uploading launch.json...")
+    launch_result = subprocess.run(
+        ["scp", *ssh_opts, *scp_port_opts, str(launch_json_path), f"{user}@{host}:/tmp/launch.json"],
+        capture_output=True,
+        text=True,
+    )
+    if launch_result.returncode != 0:
+        click.secho(f"  SCP failed: {launch_result.stderr}", fg="red", err=True)
+        return False
+
+    # SCP the VSCode tasks.json
+    tasks_json_path = script_path.parent / "tasks.json"
+    click.echo("  Uploading tasks.json...")
+    tasks_result = subprocess.run(
+        ["scp", *ssh_opts, *scp_port_opts, str(tasks_json_path), f"{user}@{host}:/tmp/tasks.json"],
+        capture_output=True,
+        text=True,
+    )
+    if tasks_result.returncode != 0:
+        click.secho(f"  SCP failed: {tasks_result.stderr}", fg="red", err=True)
+        return False
+
     # SCP the script
     click.echo("  Uploading remote_setup.sh...")
     scp_result = subprocess.run(
-        ["scp", *ssh_opts, str(script_path), f"{user}@{host}:/tmp/remote_setup.sh"],
+        ["scp", *ssh_opts, *scp_port_opts, str(script_path), f"{user}@{host}:/tmp/remote_setup.sh"],
         capture_output=True,
         text=True,
     )
@@ -156,10 +206,14 @@ def run_remote_setup(host: str, user: str, key_path: Path, script_path: Path) ->
         click.secho(f"  SCP failed: {scp_result.stderr}", fg="red", err=True)
         return False
 
-    # Execute the script
+    # Execute the script, passing PYTHON_VERSION as an inline env var if specified
     click.echo("  Running remote_setup.sh on instance...")
+    remote_cmd = "chmod +x /tmp/remote_setup.sh && "
+    if python_version:
+        remote_cmd += f"PYTHON_VERSION={python_version} "
+    remote_cmd += "/tmp/remote_setup.sh"
     ssh_result = subprocess.run(
-        ["ssh", *ssh_opts, f"{user}@{host}", "chmod +x /tmp/remote_setup.sh && /tmp/remote_setup.sh"],
+        ["ssh", *ssh_opts, *ssh_port_opts, f"{user}@{host}", remote_cmd],
         capture_output=False,
     )
     return ssh_result.returncode == 0
@@ -222,15 +276,17 @@ def _next_alias(content: str, prefix: str = "aws-gpu") -> str:
     return f"{prefix}{max_n + 1}"
 
 
-def _build_stanza(instance_id: str, alias: str, hostname: str, user: str, key_path: Path) -> str:
+def _build_stanza(instance_id: str, alias: str, hostname: str, user: str, key_path: Path, port: int = 22) -> str:
     """Build a complete SSH config stanza with markers."""
     priv_key = private_key_path(key_path)
+    port_line = f"    Port {port}\n" if port != 22 else ""
     return (
         f"{_BEGIN_MARKER.format(instance_id=instance_id)}\n"
         f"Host {alias}\n"
         f"    HostName {hostname}\n"
         f"    User {user}\n"
         f"    IdentityFile {priv_key}\n"
+        f"{port_line}"
         f"    StrictHostKeyChecking no\n"
         f"    UserKnownHostsFile /dev/null\n"
         f"{_END_MARKER.format(instance_id=instance_id)}\n"
@@ -244,6 +300,7 @@ def add_ssh_host(
     key_path: Path,
     config_path: Path | None = None,
     alias_prefix: str = "aws-gpu",
+    port: int = 22,
 ) -> str:
     """Add (or update) an SSH host stanza for *instance_id*.
 
@@ -257,7 +314,7 @@ def add_ssh_host(
     content = _remove_block(content, instance_id)
 
     alias = existing_alias or _next_alias(content, alias_prefix)
-    stanza = _build_stanza(instance_id, alias, hostname, user, key_path)
+    stanza = _build_stanza(instance_id, alias, hostname, user, key_path, port=port)
 
     # Ensure a blank line before our block if file has content
     if content and not content.endswith("\n\n") and not content.endswith("\n"):
@@ -317,21 +374,6 @@ def list_ssh_hosts(config_path: Path | None = None) -> dict[str, str]:
     return result
 
 
-# ---------------------------------------------------------------------------
-# GPU info via SSH
-# ---------------------------------------------------------------------------
-
-_GPU_ARCHITECTURES: dict[str, str] = {
-    "7.0": "Volta",
-    "7.5": "Turing",
-    "8.0": "Ampere",
-    "8.6": "Ampere",
-    "8.7": "Ampere",
-    "8.9": "Ada Lovelace",
-    "9.0": "Hopper",
-}
-
-
 @dataclass
 class SSHHostDetails:
     """Connection details parsed from an SSH config stanza."""
@@ -339,18 +381,7 @@ class SSHHostDetails:
     hostname: str
     user: str
     identity_file: Path
-
-
-@dataclass
-class GpuInfo:
-    """GPU information retrieved via nvidia-smi and nvcc."""
-
-    driver_version: str
-    cuda_driver_version: str  # max CUDA version supported by driver (from nvidia-smi)
-    cuda_toolkit_version: str | None  # actual CUDA toolkit installed (from nvcc), None if unavailable
-    gpu_name: str
-    compute_capability: str
-    architecture: str
+    port: int = 22
 
 
 def get_ssh_host_details(instance_id: str, config_path: Path | None = None) -> SSHHostDetails | None:
@@ -371,6 +402,7 @@ def get_ssh_host_details(instance_id: str, config_path: Path | None = None) -> S
     hostname: str | None = None
     user: str | None = None
     identity_file: str | None = None
+    port: int = 22
 
     for line in content.splitlines():
         if line == begin_marker:
@@ -378,7 +410,7 @@ def get_ssh_host_details(instance_id: str, config_path: Path | None = None) -> S
             continue
         if line == end_marker and in_block:
             if hostname and user and identity_file:
-                return SSHHostDetails(hostname=hostname, user=user, identity_file=Path(identity_file))
+                return SSHHostDetails(hostname=hostname, user=user, identity_file=Path(identity_file), port=port)
             return None
         if in_block:
             stripped = line.strip()
@@ -388,17 +420,20 @@ def get_ssh_host_details(instance_id: str, config_path: Path | None = None) -> S
                 user = stripped.removeprefix("User ").strip()
             elif stripped.startswith("IdentityFile "):
                 identity_file = stripped.removeprefix("IdentityFile ").strip()
+            elif stripped.startswith("Port "):
+                port = int(stripped.removeprefix("Port ").strip())
 
     return None
 
 
-def query_gpu_info(host: str, user: str, key_path: Path, timeout: int = 10) -> GpuInfo | None:
+def query_gpu_info(host: str, user: str, key_path: Path, timeout: int = 10, port: int = 22) -> GpuInfo | None:
     """SSH into a host and query GPU info via ``nvidia-smi``.
 
     Returns ``GpuInfo`` on success, or ``None`` if the SSH connection fails,
     ``nvidia-smi`` is unavailable, or the output is malformed.
     """
     ssh_opts = _ssh_opts(key_path)
+    port_opts = ["-p", str(port)] if port != 22 else []
     remote_cmd = (
         "nvidia-smi --query-gpu=driver_version,name,compute_cap --format=csv,noheader,nounits"
         " && nvidia-smi | grep -oP 'CUDA Version: \\K[\\d.]+'"
@@ -407,6 +442,7 @@ def query_gpu_info(host: str, user: str, key_path: Path, timeout: int = 10) -> G
     cmd = [
         "ssh",
         *ssh_opts,
+        *port_opts,
         "-o",
         f"ConnectTimeout={timeout}",
         "-o",

aws_bootstrap/tests/test_cli.py

@@ -9,7 +9,8 @@ import botocore.exceptions
 from click.testing import CliRunner
 
 from aws_bootstrap.cli import main
-from aws_bootstrap.ssh import GpuInfo, SSHHostDetails
+from aws_bootstrap.gpu import GpuInfo
+from aws_bootstrap.ssh import SSHHostDetails
 
 
 def test_help():
@@ -73,11 +74,12 @@ def test_status_no_instances(mock_find, mock_session):
     assert "No active" in result.output
 
 
+@patch("aws_bootstrap.cli.get_ssh_host_details", return_value=None)
 @patch("aws_bootstrap.cli.list_ssh_hosts", return_value={})
 @patch("aws_bootstrap.cli.boto3.Session")
 @patch("aws_bootstrap.cli.get_spot_price")
 @patch("aws_bootstrap.cli.find_tagged_instances")
-def test_status_shows_instances(mock_find, mock_spot_price, mock_session, mock_ssh_hosts):
+def test_status_shows_instances(mock_find, mock_spot_price, mock_session, mock_ssh_hosts, mock_details):
     mock_find.return_value = [
         {
             "InstanceId": "i-abc123",
@@ -101,11 +103,12 @@ def test_status_shows_instances(mock_find, mock_spot_price, mock_session, mock_s
     assert "Est. cost" in result.output
 
 
+@patch("aws_bootstrap.cli.get_ssh_host_details", return_value=None)
 @patch("aws_bootstrap.cli.list_ssh_hosts", return_value={})
 @patch("aws_bootstrap.cli.boto3.Session")
 @patch("aws_bootstrap.cli.get_spot_price")
 @patch("aws_bootstrap.cli.find_tagged_instances")
-def test_status_on_demand_no_cost(mock_find, mock_spot_price, mock_session, mock_ssh_hosts):
+def test_status_on_demand_no_cost(mock_find, mock_spot_price, mock_session, mock_ssh_hosts, mock_details):
     mock_find.return_value = [
         {
             "InstanceId": "i-ondemand",
@@ -351,11 +354,12 @@ def test_terminate_removes_ssh_config(mock_terminate, mock_find, mock_session, m
     mock_remove_ssh.assert_called_once_with("i-abc123")
 
 
+@patch("aws_bootstrap.cli.get_ssh_host_details", return_value=None)
 @patch("aws_bootstrap.cli.list_ssh_hosts")
 @patch("aws_bootstrap.cli.boto3.Session")
 @patch("aws_bootstrap.cli.get_spot_price")
 @patch("aws_bootstrap.cli.find_tagged_instances")
-def test_status_shows_alias(mock_find, mock_spot_price, mock_session, mock_ssh_hosts):
+def test_status_shows_alias(mock_find, mock_spot_price, mock_session, mock_ssh_hosts, mock_details):
     mock_find.return_value = [
         {
             "InstanceId": "i-abc123",
@@ -376,11 +380,12 @@ def test_status_shows_alias(mock_find, mock_spot_price, mock_session, mock_ssh_h
     assert "aws-gpu1" in result.output
 
 
+@patch("aws_bootstrap.cli.get_ssh_host_details", return_value=None)
 @patch("aws_bootstrap.cli.list_ssh_hosts", return_value={})
 @patch("aws_bootstrap.cli.boto3.Session")
 @patch("aws_bootstrap.cli.get_spot_price")
 @patch("aws_bootstrap.cli.find_tagged_instances")
-def test_status_no_alias_graceful(mock_find, mock_spot_price, mock_session, mock_ssh_hosts):
+def test_status_no_alias_graceful(mock_find, mock_spot_price, mock_session, mock_ssh_hosts, mock_details):
     mock_find.return_value = [
         {
             "InstanceId": "i-old999",
@@ -520,13 +525,98 @@ def test_status_gpu_skips_non_running(mock_find, mock_session, mock_ssh_hosts, m
 @patch("aws_bootstrap.cli.boto3.Session")
 @patch("aws_bootstrap.cli.get_spot_price", return_value=0.15)
 @patch("aws_bootstrap.cli.find_tagged_instances")
-def test_status_without_gpu_flag_no_ssh(mock_find, mock_spot, mock_session, mock_ssh_hosts, mock_details, mock_gpu):
+def test_status_without_gpu_flag_no_gpu_query(
+    mock_find, mock_spot, mock_session, mock_ssh_hosts, mock_details, mock_gpu
+):
     mock_find.return_value = [_RUNNING_INSTANCE]
     runner = CliRunner()
     result = runner.invoke(main, ["status"])
     assert result.exit_code == 0
     mock_gpu.assert_not_called()
-    mock_details.assert_not_called()
+
+
+# ---------------------------------------------------------------------------
+# --instructions / --no-instructions / -I flag tests
+# ---------------------------------------------------------------------------
+
+
+def test_status_help_shows_instructions_flag():
+    runner = CliRunner()
+    result = runner.invoke(main, ["status", "--help"])
+    assert result.exit_code == 0
+    assert "--instructions" in result.output
+    assert "--no-instructions" in result.output
+    assert "-I" in result.output
+
+
+@patch("aws_bootstrap.cli.get_ssh_host_details")
+@patch("aws_bootstrap.cli.list_ssh_hosts", return_value={"i-abc123": "aws-gpu1"})
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.get_spot_price", return_value=0.15)
+@patch("aws_bootstrap.cli.find_tagged_instances")
+def test_status_instructions_shown_by_default(mock_find, mock_spot, mock_session, mock_ssh_hosts, mock_details):
+    """Instructions are shown by default (no flag needed)."""
+    mock_find.return_value = [_RUNNING_INSTANCE]
+    mock_details.return_value = SSHHostDetails(
+        hostname="1.2.3.4", user="ubuntu", identity_file=Path("/home/user/.ssh/id_ed25519")
+    )
+    runner = CliRunner()
+    result = runner.invoke(main, ["status"])
+    assert result.exit_code == 0
+    assert "ssh aws-gpu1" in result.output
+    assert "ssh -NL 8888:localhost:8888 aws-gpu1" in result.output
+    assert "vscode-remote://ssh-remote+aws-gpu1/home/ubuntu/workspace" in result.output
+    assert "python ~/gpu_benchmark.py" in result.output
+
+
+@patch("aws_bootstrap.cli.get_ssh_host_details")
+@patch("aws_bootstrap.cli.list_ssh_hosts", return_value={"i-abc123": "aws-gpu1"})
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.get_spot_price", return_value=0.15)
+@patch("aws_bootstrap.cli.find_tagged_instances")
+def test_status_no_instructions_suppresses_commands(mock_find, mock_spot, mock_session, mock_ssh_hosts, mock_details):
+    """--no-instructions suppresses connection commands."""
+    mock_find.return_value = [_RUNNING_INSTANCE]
+    mock_details.return_value = SSHHostDetails(
+        hostname="1.2.3.4", user="ubuntu", identity_file=Path("/home/user/.ssh/id_ed25519")
+    )
+    runner = CliRunner()
+    result = runner.invoke(main, ["status", "--no-instructions"])
+    assert result.exit_code == 0
+    assert "vscode-remote" not in result.output
+    assert "Jupyter" not in result.output
+
+
+@patch("aws_bootstrap.cli.get_ssh_host_details")
+@patch("aws_bootstrap.cli.list_ssh_hosts", return_value={})
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.get_spot_price", return_value=0.15)
+@patch("aws_bootstrap.cli.find_tagged_instances")
+def test_status_instructions_no_alias_skips(mock_find, mock_spot, mock_session, mock_ssh_hosts, mock_details):
+    """Instances without an SSH alias don't get connection instructions."""
+    mock_find.return_value = [_RUNNING_INSTANCE]
+    runner = CliRunner()
+    result = runner.invoke(main, ["status"])
+    assert result.exit_code == 0
+    assert "ssh aws-gpu" not in result.output
+    assert "vscode-remote" not in result.output
+
+
+@patch("aws_bootstrap.cli.get_ssh_host_details")
+@patch("aws_bootstrap.cli.list_ssh_hosts", return_value={"i-abc123": "aws-gpu1"})
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.get_spot_price", return_value=0.15)
+@patch("aws_bootstrap.cli.find_tagged_instances")
+def test_status_instructions_non_default_port(mock_find, mock_spot, mock_session, mock_ssh_hosts, mock_details):
+    mock_find.return_value = [_RUNNING_INSTANCE]
+    mock_details.return_value = SSHHostDetails(
+        hostname="1.2.3.4", user="ubuntu", identity_file=Path("/home/user/.ssh/id_ed25519"), port=2222
+    )
+    runner = CliRunner()
+    result = runner.invoke(main, ["status"])
+    assert result.exit_code == 0
+    assert "ssh -p 2222 aws-gpu1" in result.output
+    assert "ssh -NL 8888:localhost:8888 -p 2222 aws-gpu1" in result.output
 
 
 # ---------------------------------------------------------------------------
@@ -636,3 +726,111 @@ def test_no_credentials_caught_on_list(mock_session, mock_list):
     result = runner.invoke(main, ["list", "instance-types"])
     assert result.exit_code != 0
     assert "Unable to locate AWS credentials" in result.output
+
+
+# ---------------------------------------------------------------------------
+# --python-version tests
+# ---------------------------------------------------------------------------
+
+
+@patch("aws_bootstrap.cli.add_ssh_host", return_value="aws-gpu1")
+@patch("aws_bootstrap.cli.run_remote_setup", return_value=True)
+@patch("aws_bootstrap.cli.wait_for_ssh", return_value=True)
+@patch("aws_bootstrap.cli.wait_instance_ready")
+@patch("aws_bootstrap.cli.launch_instance")
+@patch("aws_bootstrap.cli.ensure_security_group", return_value="sg-123")
+@patch("aws_bootstrap.cli.import_key_pair", return_value="aws-bootstrap-key")
+@patch("aws_bootstrap.cli.get_latest_ami")
+@patch("aws_bootstrap.cli.boto3.Session")
+def test_launch_python_version_passed_to_setup(
+    mock_session, mock_ami, mock_import, mock_sg, mock_launch, mock_wait, mock_ssh, mock_setup, mock_add_ssh, tmp_path
+):
+    mock_ami.return_value = {"ImageId": "ami-123", "Name": "TestAMI"}
+    mock_launch.return_value = {"InstanceId": "i-test123"}
+    mock_wait.return_value = {"PublicIpAddress": "1.2.3.4"}
+
+    key_path = tmp_path / "id_ed25519.pub"
+    key_path.write_text("ssh-ed25519 AAAA test@host")
+
+    runner = CliRunner()
+    result = runner.invoke(main, ["launch", "--key-path", str(key_path), "--python-version", "3.13"])
+    assert result.exit_code == 0
+    mock_setup.assert_called_once()
+    assert mock_setup.call_args[0][4] == "3.13"
+
+
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.get_latest_ami")
+@patch("aws_bootstrap.cli.import_key_pair", return_value="aws-bootstrap-key")
+@patch("aws_bootstrap.cli.ensure_security_group", return_value="sg-123")
+def test_launch_dry_run_shows_python_version(mock_sg, mock_import, mock_ami, mock_session, tmp_path):
+    mock_ami.return_value = {"ImageId": "ami-123", "Name": "TestAMI"}
+
+    key_path = tmp_path / "id_ed25519.pub"
+    key_path.write_text("ssh-ed25519 AAAA test@host")
+
+    runner = CliRunner()
+    result = runner.invoke(main, ["launch", "--key-path", str(key_path), "--dry-run", "--python-version", "3.14.2"])
+    assert result.exit_code == 0
+    assert "3.14.2" in result.output
+    assert "Python version" in result.output
+
+
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.get_latest_ami")
+@patch("aws_bootstrap.cli.import_key_pair", return_value="aws-bootstrap-key")
+@patch("aws_bootstrap.cli.ensure_security_group", return_value="sg-123")
+def test_launch_dry_run_omits_python_version_when_unset(mock_sg, mock_import, mock_ami, mock_session, tmp_path):
+    mock_ami.return_value = {"ImageId": "ami-123", "Name": "TestAMI"}
+
+    key_path = tmp_path / "id_ed25519.pub"
+    key_path.write_text("ssh-ed25519 AAAA test@host")
+
+    runner = CliRunner()
+    result = runner.invoke(main, ["launch", "--key-path", str(key_path), "--dry-run"])
+    assert result.exit_code == 0
+    assert "Python version" not in result.output
+
+
+# ---------------------------------------------------------------------------
+# --ssh-port tests
+# ---------------------------------------------------------------------------
+
+
+def test_launch_help_shows_ssh_port():
+    runner = CliRunner()
+    result = runner.invoke(main, ["launch", "--help"])
+    assert result.exit_code == 0
+    assert "--ssh-port" in result.output
+
+
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.get_latest_ami")
+@patch("aws_bootstrap.cli.import_key_pair", return_value="aws-bootstrap-key")
+@patch("aws_bootstrap.cli.ensure_security_group", return_value="sg-123")
+def test_launch_dry_run_shows_ssh_port_when_non_default(mock_sg, mock_import, mock_ami, mock_session, tmp_path):
+    mock_ami.return_value = {"ImageId": "ami-123", "Name": "TestAMI"}
+
+    key_path = tmp_path / "id_ed25519.pub"
+    key_path.write_text("ssh-ed25519 AAAA test@host")
+
+    runner = CliRunner()
+    result = runner.invoke(main, ["launch", "--key-path", str(key_path), "--dry-run", "--ssh-port", "2222"])
+    assert result.exit_code == 0
+    assert "2222" in result.output
+
+
+@patch("aws_bootstrap.cli.boto3.Session")
+@patch("aws_bootstrap.cli.get_latest_ami")
+@patch("aws_bootstrap.cli.import_key_pair", return_value="aws-bootstrap-key")
+@patch("aws_bootstrap.cli.ensure_security_group", return_value="sg-123")
+def test_launch_dry_run_omits_ssh_port_when_default(mock_sg, mock_import, mock_ami, mock_session, tmp_path):
+    mock_ami.return_value = {"ImageId": "ami-123", "Name": "TestAMI"}
+
+    key_path = tmp_path / "id_ed25519.pub"
+    key_path.write_text("ssh-ed25519 AAAA test@host")
+
+    runner = CliRunner()
+    result = runner.invoke(main, ["launch", "--key-path", str(key_path), "--dry-run"])
+    assert result.exit_code == 0
+    assert "SSH port" not in result.output

aws_bootstrap/tests/test_gpu.py

@@ -0,0 +1,98 @@
+"""Tests for GPU info queries via SSH (query_gpu_info, GPU architecture mapping)."""
+
+from __future__ import annotations
+import subprocess
+from pathlib import Path
+from unittest.mock import patch
+
+from aws_bootstrap.gpu import _GPU_ARCHITECTURES, GpuInfo
+from aws_bootstrap.ssh import query_gpu_info
+
+
+# ---------------------------------------------------------------------------
+# query_gpu_info
+# ---------------------------------------------------------------------------
+
+NVIDIA_SMI_OUTPUT = "560.35.03, Tesla T4, 7.5\n12.8\n12.6\n"
+
+
+@patch("aws_bootstrap.ssh.subprocess.run")
+def test_query_gpu_info_success(mock_run):
+    """Successful nvidia-smi + nvcc output returns a valid GpuInfo."""
+    mock_run.return_value = subprocess.CompletedProcess(args=[], returncode=0, stdout=NVIDIA_SMI_OUTPUT, stderr="")
+
+    info = query_gpu_info("1.2.3.4", "ubuntu", Path("/home/user/.ssh/id_ed25519"))
+    assert info is not None
+    assert isinstance(info, GpuInfo)
+    assert info.driver_version == "560.35.03"
+    assert info.cuda_driver_version == "12.8"
+    assert info.cuda_toolkit_version == "12.6"
+    assert info.gpu_name == "Tesla T4"
+    assert info.compute_capability == "7.5"
+    assert info.architecture == "Turing"
+
+
+@patch("aws_bootstrap.ssh.subprocess.run")
+def test_query_gpu_info_no_nvcc(mock_run):
+    """When nvcc is unavailable, cuda_toolkit_version is None."""
+    output = "560.35.03, Tesla T4, 7.5\n12.8\nN/A\n"
+    mock_run.return_value = subprocess.CompletedProcess(args=[], returncode=0, stdout=output, stderr="")
+
+    info = query_gpu_info("1.2.3.4", "ubuntu", Path("/home/user/.ssh/id_ed25519"))
+    assert info is not None
+    assert info.cuda_driver_version == "12.8"
+    assert info.cuda_toolkit_version is None
+
+
+@patch("aws_bootstrap.ssh.subprocess.run")
+def test_query_gpu_info_ssh_failure(mock_run):
+    """Non-zero exit code returns None."""
+    mock_run.return_value = subprocess.CompletedProcess(args=[], returncode=255, stdout="", stderr="Connection refused")
+
+    info = query_gpu_info("1.2.3.4", "ubuntu", Path("/home/user/.ssh/id_ed25519"))
+    assert info is None
+
+
+@patch("aws_bootstrap.ssh.subprocess.run", side_effect=subprocess.TimeoutExpired(cmd="ssh", timeout=15))
+def test_query_gpu_info_timeout(mock_run):
+    """TimeoutExpired returns None."""
+    info = query_gpu_info("1.2.3.4", "ubuntu", Path("/home/user/.ssh/id_ed25519"))
+    assert info is None
+
+
+@patch("aws_bootstrap.ssh.subprocess.run")
+def test_query_gpu_info_malformed_output(mock_run):
+    """Garbage output returns None."""
+    mock_run.return_value = subprocess.CompletedProcess(
+        args=[], returncode=0, stdout="not valid gpu output\n", stderr=""
+    )
+
+    info = query_gpu_info("1.2.3.4", "ubuntu", Path("/home/user/.ssh/id_ed25519"))
+    assert info is None
+
+
+# ---------------------------------------------------------------------------
+# GPU architecture mapping
+# ---------------------------------------------------------------------------
+
+
+def test_gpu_architecture_mapping():
+    """Known compute capabilities map to correct architecture names."""
+    assert _GPU_ARCHITECTURES["7.5"] == "Turing"
+    assert _GPU_ARCHITECTURES["8.0"] == "Ampere"
+    assert _GPU_ARCHITECTURES["8.6"] == "Ampere"
+    assert _GPU_ARCHITECTURES["8.9"] == "Ada Lovelace"
+    assert _GPU_ARCHITECTURES["9.0"] == "Hopper"
+    assert _GPU_ARCHITECTURES["7.0"] == "Volta"
+
+
+@patch("aws_bootstrap.ssh.subprocess.run")
+def test_query_gpu_info_unknown_architecture(mock_run):
+    """Unknown compute capability produces a fallback architecture string."""
+    mock_run.return_value = subprocess.CompletedProcess(
+        args=[], returncode=0, stdout="550.00.00, Future GPU, 10.0\n13.0\n13.0\n", stderr=""
+    )
+
+    info = query_gpu_info("1.2.3.4", "ubuntu", Path("/home/user/.ssh/id_ed25519"))
+    assert info is not None
+    assert info.architecture == "Unknown (10.0)"