aws-bootstrap-g4dn 0.1.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

aws_bootstrap/cli.py

@@ -5,6 +5,7 @@ from datetime import UTC, datetime
 from pathlib import Path
 
 import boto3
+import botocore.exceptions
 import click
 
 from .config import LaunchConfig
@@ -56,7 +57,39 @@ def warn(msg: str) -> None:
     click.secho(f"  WARNING: {msg}", fg="yellow", err=True)
 
 
-@click.group()
+class _AWSGroup(click.Group):
+    """Click group that catches common AWS credential/auth errors."""
+
+    def invoke(self, ctx):
+        try:
+            return super().invoke(ctx)
+        except botocore.exceptions.NoCredentialsError:
+            raise CLIError(
+                "Unable to locate AWS credentials.\n\n"
+                "  Make sure you have configured AWS credentials using one of:\n"
+                "    - Set the AWS_PROFILE environment variable:  export AWS_PROFILE=<profile-name>\n"
+                "    - Pass --profile to the command:  aws-bootstrap <command> --profile <profile-name>\n"
+                "    - Configure a default profile:  aws configure\n\n"
+                "  See: https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-files.html"
+            ) from None
+        except botocore.exceptions.ProfileNotFound as e:
+            raise CLIError(f"{e}\n\n  List available profiles with:  aws configure list-profiles") from None
+        except botocore.exceptions.PartialCredentialsError as e:
+            raise CLIError(
+                f"Incomplete AWS credentials: {e}\n\n  Check your AWS configuration with:  aws configure list"
+            ) from None
+        except botocore.exceptions.ClientError as e:
+            code = e.response["Error"]["Code"]
+            if code in ("AuthFailure", "UnauthorizedOperation", "ExpiredTokenException", "ExpiredToken"):
+                raise CLIError(
+                    f"AWS authorization failed: {e.response['Error']['Message']}\n\n"
+                    "  Your credentials may be expired or lack the required permissions.\n"
+                    "  Check your AWS configuration with:  aws configure list"
+                ) from None
+            raise
+
+
+@click.group(cls=_AWSGroup)
 @click.version_option(package_name="aws-bootstrap-g4dn")
 def main():
     """Bootstrap AWS EC2 GPU instances for hybrid local-remote development."""
@@ -80,6 +113,12 @@ def main():
 @click.option("--no-setup", is_flag=True, default=False, help="Skip running the remote setup script.")
 @click.option("--dry-run", is_flag=True, default=False, help="Show what would be done without executing.")
 @click.option("--profile", default=None, help="AWS profile override (defaults to AWS_PROFILE env var).")
+@click.option(
+    "--python-version",
+    default=None,
+    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.")
 def launch(
     instance_type,
     ami_filter,
@@ -92,6 +131,8 @@ def launch(
     no_setup,
     dry_run,
     profile,
+    python_version,
+    ssh_port,
 ):
     """Launch a GPU-accelerated EC2 instance."""
     config = LaunchConfig(
@@ -104,6 +145,8 @@ def launch(
         volume_size=volume_size,
         run_setup=not no_setup,
         dry_run=dry_run,
+        ssh_port=ssh_port,
+        python_version=python_version,
     )
     if ami_filter:
         config.ami_filter = ami_filter
@@ -130,7 +173,7 @@ def launch(
 
     # Step 3: Security group
     step(3, 6, "Ensuring security group...")
-    sg_id = ensure_security_group(ec2, config.security_group, config.tag_value)
+    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"
 
@@ -145,6 +188,10 @@ def launch(
         val("Volume", f"{config.volume_size} GB gp3")
         val("Region", config.region)
         val("Remote setup", "yes" if config.run_setup else "no")
+        if config.ssh_port != 22:
+            val("SSH port", str(config.ssh_port))
+        if config.python_version:
+            val("Python version", config.python_version)
         click.echo()
         click.secho("No resources launched (dry-run mode).", fg="yellow")
         return
@@ -169,9 +216,13 @@ def launch(
     # Step 6: SSH and remote setup
     step(6, 6, "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):
+    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.")
-        info(f"Instance is running — try connecting manually: ssh -i {private_key} {config.ssh_user}@{public_ip}")
+        port_flag = f" -p {config.ssh_port}" if config.ssh_port != 22 else ""
+        info(
+            f"Instance is running — try connecting manually:"
+            f" ssh -i {private_key}{port_flag} {config.ssh_user}@{public_ip}"
+        )
         return
 
     if config.run_setup:
@@ -179,7 +230,9 @@ def launch(
             warn(f"Setup script not found at {SETUP_SCRIPT}, skipping.")
         else:
             info("Running remote setup...")
-            if run_remote_setup(public_ip, config.ssh_user, config.key_path, SETUP_SCRIPT):
+            if run_remote_setup(
+                public_ip, config.ssh_user, config.key_path, SETUP_SCRIPT, config.python_version, port=config.ssh_port
+            ):
                 success("Remote setup completed successfully.")
             else:
                 warn("Remote setup failed. Instance is still running.")
@@ -191,6 +244,7 @@ def launch(
         user=config.ssh_user,
         key_path=config.key_path,
         alias_prefix=config.alias_prefix,
+        port=config.ssh_port,
     )
     success(f"Added SSH config alias: {alias}")
 
@@ -206,18 +260,27 @@ def launch(
     val("Pricing", pricing)
     val("SSH alias", alias)
 
+    port_flag = f" -p {config.ssh_port}" if config.ssh_port != 22 else ""
+
     click.echo()
     click.secho("  SSH:", fg="cyan")
-    click.secho(f"    ssh {alias}", bold=True)
-    info(f"or: ssh -i {private_key} {config.ssh_user}@{public_ip}")
+    click.secho(f"    ssh{port_flag} {alias}", bold=True)
+    info(f"or: ssh -i {private_key}{port_flag} {config.ssh_user}@{public_ip}")
 
     click.echo()
     click.secho("  Jupyter (via SSH tunnel):", fg="cyan")
-    click.secho(f"    ssh -NL 8888:localhost:8888 {alias}", bold=True)
-    info(f"or: ssh -i {private_key} -NL 8888:localhost:8888 {config.ssh_user}@{public_ip}")
+    click.secho(f"    ssh -NL 8888:localhost:8888{port_flag} {alias}", bold=True)
+    info(f"or: ssh -i {private_key} -NL 8888:localhost:8888{port_flag} {config.ssh_user}@{public_ip}")
     info("Then open: http://localhost:8888")
     info("Notebook: ~/gpu_smoke_test.ipynb (GPU smoke test)")
 
+    click.echo()
+    click.secho("  VSCode Remote SSH:", fg="cyan")
+    click.secho(
+        f"    code --folder-uri vscode-remote://ssh-remote+{alias}/home/{config.ssh_user}",
+        bold=True,
+    )
+
     click.echo()
     click.secho("  GPU Benchmark:", fg="cyan")
     click.secho(f"    ssh {alias} 'python ~/gpu_benchmark.py'", bold=True)
@@ -233,7 +296,14 @@ def launch(
 @click.option("--region", default="us-west-2", show_default=True, help="AWS region.")
 @click.option("--profile", default=None, help="AWS profile override.")
 @click.option("--gpu", is_flag=True, default=False, help="Query GPU info (CUDA, driver) via SSH.")
-def status(region, profile, gpu):
+@click.option(
+    "--instructions/--no-instructions",
+    "-I",
+    default=True,
+    show_default=True,
+    help="Show connection commands (SSH, Jupyter, VSCode) for each running instance.",
+)
+def status(region, profile, gpu, instructions):
     """Show running instances created by aws-bootstrap."""
     session = boto3.Session(profile_name=profile, region_name=region)
     ec2 = session.client("ec2")
@@ -272,11 +342,15 @@ def status(region, profile, gpu):
         if inst["PublicIp"]:
             val("    IP", inst["PublicIp"])
 
+        # Look up SSH config details once (used by --gpu and --with-instructions)
+        details = None
+        if (gpu or instructions) and state == "running" and inst["PublicIp"]:
+            details = get_ssh_host_details(inst["InstanceId"])
+
         # GPU info (opt-in, only for running instances with a public IP)
         if gpu and state == "running" and inst["PublicIp"]:
-            details = get_ssh_host_details(inst["InstanceId"])
             if details:
-                gpu_info = query_gpu_info(details.hostname, details.user, details.identity_file)
+                gpu_info = query_gpu_info(details.hostname, details.user, details.identity_file, port=details.port)
             else:
                 gpu_info = query_gpu_info(
                     inst["PublicIp"],
@@ -320,6 +394,29 @@ def status(region, profile, gpu):
                 val("    Est. cost", f"~${est_cost:.4f}")
 
         val("    Launched", str(inst["LaunchTime"]))
+
+        # Connection instructions (opt-in, only for running instances with a public IP and alias)
+        if instructions and state == "running" and inst["PublicIp"] and alias:
+            user = details.user if details else "ubuntu"
+            port = details.port if details else 22
+            port_flag = f" -p {port}" if port != 22 else ""
+
+            click.echo()
+            click.secho("    SSH:", fg="cyan")
+            click.secho(f"      ssh{port_flag} {alias}", bold=True)
+
+            click.secho("    Jupyter (via SSH tunnel):", fg="cyan")
+            click.secho(f"      ssh -NL 8888:localhost:8888{port_flag} {alias}", bold=True)
+
+            click.secho("    VSCode Remote SSH:", fg="cyan")
+            click.secho(
+                f"      code --folder-uri vscode-remote://ssh-remote+{alias}/home/{user}",
+                bold=True,
+            )
+
+            click.secho("    GPU Benchmark:", fg="cyan")
+            click.secho(f"      ssh {alias} 'python ~/gpu_benchmark.py'", bold=True)
+
     click.echo()
     first_id = instances[0]["InstanceId"]
     click.echo("  To terminate:  " + click.style(f"aws-bootstrap terminate {first_id}", bold=True))

aws_bootstrap/config.py

@@ -22,3 +22,5 @@ class LaunchConfig:
     ssh_user: str = "ubuntu"
     tag_value: str = "aws-bootstrap-g4dn"
     alias_prefix: str = "aws-gpu"
+    ssh_port: int = 22
+    python_version: str | None = None

aws_bootstrap/ec2.py

@@ -59,7 +59,7 @@ def get_latest_ami(ec2_client, ami_filter: str) -> dict:
     return images[0]
 
 
-def ensure_security_group(ec2_client, name: str, tag_value: str) -> str:
+def ensure_security_group(ec2_client, name: str, tag_value: str, ssh_port: int = 22) -> str:
     """Find or create a security group with SSH ingress in the default VPC."""
     # Find default VPC
     vpcs = ec2_client.describe_vpcs(Filters=[{"Name": "isDefault", "Values": ["true"]}])
@@ -103,8 +103,8 @@ def ensure_security_group(ec2_client, name: str, tag_value: str) -> str:
         IpPermissions=[
             {
                 "IpProtocol": "tcp",
-                "FromPort": 22,
-                "ToPort": 22,
+                "FromPort": ssh_port,
+                "ToPort": ssh_port,
                 "IpRanges": [{"CidrIp": "0.0.0.0/0", "Description": "SSH access"}],
             }
         ],

aws_bootstrap/gpu.py

@@ -0,0 +1,27 @@
+"""GPU architecture mapping and GPU info dataclass."""
+
+from __future__ import annotations
+from dataclasses import dataclass
+
+
+_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 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

aws_bootstrap/resources/remote_setup.sh

@@ -34,7 +34,13 @@ if ! command -v uv &>/dev/null; then
 fi
 export PATH="$HOME/.local/bin:$PATH"
 
-uv venv ~/venv
+if [ -n "${PYTHON_VERSION:-}" ]; then
+    echo "  Installing Python ${PYTHON_VERSION}..."
+    uv python install "$PYTHON_VERSION"
+    uv venv --python "$PYTHON_VERSION" ~/venv
+else
+    uv venv ~/venv
+fi
 
 # --- CUDA-aware PyTorch installation ---
 # Known PyTorch CUDA wheel tags (ascending order).

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,
     )
@@ -148,7 +162,7 @@ def run_remote_setup(host: str, user: str, key_path: Path, script_path: Path) ->
     # 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 +170,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 +240,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 +264,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 +278,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 +338,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 +345,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 +366,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 +374,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 +384,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 +406,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",