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

aws_bootstrap/cli.py

@@ -113,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,
@@ -125,6 +131,8 @@ def launch(
     no_setup,
     dry_run,
     profile,
+    python_version,
+    ssh_port,
 ):
     """Launch a GPU-accelerated EC2 instance."""
     config = LaunchConfig(
@@ -137,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
@@ -163,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"
 
@@ -178,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
@@ -202,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:
@@ -212,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.")
@@ -224,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}")
 
@@ -239,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}/workspace",
+        bold=True,
+    )
+
     click.echo()
     click.secho("  GPU Benchmark:", fg="cyan")
     click.secho(f"    ssh {alias} 'python ~/gpu_benchmark.py'", bold=True)
@@ -266,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")
@@ -305,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"],
@@ -353,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}/workspace",
+                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/gpu_benchmark.py

@@ -628,7 +628,9 @@ def configure_precision(device: torch.device, requested: PrecisionMode) -> Preci
     return PrecisionMode.FP32
 
 
-def print_system_info(requested_precision: PrecisionMode) -> tuple[torch.device, PrecisionMode]:
+def print_system_info(
+    requested_precision: PrecisionMode, force_cpu: bool = False
+) -> tuple[torch.device, PrecisionMode]:
     """Print system and CUDA information, return device and actual precision mode."""
     print("\n" + "=" * 60)
     print("System Information")
@@ -636,7 +638,7 @@ def print_system_info(requested_precision: PrecisionMode) -> tuple[torch.device,
     print(f"PyTorch version: {torch.__version__}")
     print(f"Python version: {sys.version.split()[0]}")
 
-    if torch.cuda.is_available():
+    if torch.cuda.is_available() and not force_cpu:
         device = torch.device("cuda")
         print("CUDA available: Yes")
         print(f"CUDA version: {torch.version.cuda}")
@@ -666,8 +668,11 @@ def print_system_info(requested_precision: PrecisionMode) -> tuple[torch.device,
     else:
         device = torch.device("cpu")
         actual_precision = PrecisionMode.FP32
-        print("CUDA available: No (running on CPU)")
-        print("WARNING: GPU benchmark results will not be representative!")
+        if force_cpu:
+            print("CPU-only mode requested (--cpu flag)")
+        else:
+            print("CUDA available: No (running on CPU)")
+        print("Running on CPU for benchmarking")
 
     print("=" * 60)
     return device, actual_precision
@@ -724,10 +729,15 @@ def main() -> None:
         action="store_true",
         help="Run CUDA/cuBLAS diagnostic tests before benchmarking",
     )
+    parser.add_argument(
+        "--cpu",
+        action="store_true",
+        help="Force CPU-only execution (for CPU vs GPU comparison)",
+    )
     args = parser.parse_args()
 
     requested_precision = PrecisionMode(args.precision)
-    device, actual_precision = print_system_info(requested_precision)
+    device, actual_precision = print_system_info(requested_precision, force_cpu=args.cpu)
 
     # Run diagnostics if requested
     if args.diagnose:

aws_bootstrap/resources/launch.json

@@ -0,0 +1,42 @@
+{
+    // CUDA debug configurations for VSCode
+    // Deployed to: ~/workspace/.vscode/launch.json
+    //
+    // Usage: Open any .cu file, press F5 to build and debug
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "CUDA: Build and Debug Active File",
+            "type": "cuda-gdb",
+            "request": "launch",
+            "program": "${fileDirname}/${fileBasenameNoExtension}",
+            "args": [],
+            "cwd": "${fileDirname}",
+            "miDebuggerPath": "__CUDA_GDB_PATH__",
+            "stopAtEntry": false,
+            "preLaunchTask": "nvcc: build active file (debug)"
+        },
+        {
+            "name": "CUDA: Build and Debug (stop at main)",
+            "type": "cuda-gdb",
+            "request": "launch",
+            "program": "${fileDirname}/${fileBasenameNoExtension}",
+            "args": [],
+            "cwd": "${fileDirname}",
+            "miDebuggerPath": "__CUDA_GDB_PATH__",
+            "stopAtEntry": true,
+            "preLaunchTask": "nvcc: build active file (debug)"
+        },
+        {
+            "name": "CUDA: Run Active File (no debug)",
+            "type": "cuda-gdb",
+            "request": "launch",
+            "program": "${fileDirname}/${fileBasenameNoExtension}",
+            "args": [],
+            "cwd": "${fileDirname}",
+            "miDebuggerPath": "__CUDA_GDB_PATH__",
+            "stopAtEntry": false,
+            "preLaunchTask": "nvcc: build active file (release)"
+        }
+    ]
+}

aws_bootstrap/resources/remote_setup.sh

@@ -7,7 +7,7 @@ echo "=== aws-bootstrap-g4dn remote setup ==="
 
 # 1. Verify GPU
 echo ""
-echo "[1/5] Verifying GPU and CUDA..."
+echo "[1/6] Verifying GPU and CUDA..."
 if command -v nvidia-smi &>/dev/null; then
     nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv,noheader
 else
@@ -20,21 +20,52 @@ else
     echo "WARNING: nvcc not found (CUDA toolkit may not be installed)"
 fi
 
+# Make Nsight Systems (nsys) available on PATH if installed under /opt/nvidia
+if ! command -v nsys &>/dev/null; then
+    NSIGHT_DIR="/opt/nvidia/nsight-systems"
+    if [ -d "$NSIGHT_DIR" ]; then
+        # Fix permissions — the parent dir is often root-only (drwx------)
+        sudo chmod o+rx "$NSIGHT_DIR"
+        # Find the latest version directory (lexicographic sort)
+        NSYS_VERSION=$(ls -1 "$NSIGHT_DIR" | sort -V | tail -1)
+        if [ -n "$NSYS_VERSION" ] && [ -x "$NSIGHT_DIR/$NSYS_VERSION/bin/nsys" ]; then
+            NSYS_BIN="$NSIGHT_DIR/$NSYS_VERSION/bin"
+            if ! grep -q "nsight-systems" ~/.bashrc 2>/dev/null; then
+                echo "export PATH=\"$NSYS_BIN:\$PATH\"" >> ~/.bashrc
+            fi
+            export PATH="$NSYS_BIN:$PATH"
+            echo "  Nsight Systems $NSYS_VERSION added to PATH ($NSYS_BIN)"
+        else
+            echo "  WARNING: Nsight Systems directory found but no nsys binary"
+        fi
+    else
+        echo "  Nsight Systems not found at $NSIGHT_DIR"
+    fi
+else
+    echo "  nsys already on PATH: $(command -v nsys)"
+fi
+
 # 2. Install utilities
 echo ""
-echo "[2/5] Installing utilities..."
+echo "[2/6] Installing utilities..."
 sudo apt-get update -qq
 sudo apt-get install -y -qq htop tmux tree jq
 
 # 3. Set up Python environment with uv
 echo ""
-echo "[3/5] Setting up Python environment with uv..."
+echo "[3/6] Setting up Python environment with uv..."
 if ! command -v uv &>/dev/null; then
     curl -LsSf https://astral.sh/uv/install.sh | sh
 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).
@@ -147,7 +178,7 @@ echo "  Jupyter config written to $JUPYTER_CONFIG_DIR/jupyter_lab_config.py"
 
 # 4. Jupyter systemd service
 echo ""
-echo "[4/5] Setting up Jupyter systemd service..."
+echo "[4/6] Setting up Jupyter systemd service..."
 LOGIN_USER=$(whoami)
 
 sudo tee /etc/systemd/system/jupyter.service > /dev/null << SVCEOF
@@ -174,7 +205,7 @@ echo "  Jupyter service started (port 8888)"
 
 # 5. SSH keepalive
 echo ""
-echo "[5/5] Configuring SSH keepalive..."
+echo "[5/6] Configuring SSH keepalive..."
 if ! grep -q "ClientAliveInterval" /etc/ssh/sshd_config; then
     echo "ClientAliveInterval 60" | sudo tee -a /etc/ssh/sshd_config > /dev/null
     echo "ClientAliveCountMax 10" | sudo tee -a /etc/ssh/sshd_config > /dev/null
@@ -184,5 +215,58 @@ else
     echo "  SSH keepalive already configured"
 fi
 
+# 6. VSCode workspace setup
+echo ""
+echo "[6/6] Setting up VSCode workspace..."
+mkdir -p ~/workspace/.vscode
+
+# Detect cuda-gdb path
+CUDA_GDB_PATH=""
+if command -v cuda-gdb &>/dev/null; then
+    CUDA_GDB_PATH=$(command -v cuda-gdb)
+elif [ -x /usr/local/cuda/bin/cuda-gdb ]; then
+    CUDA_GDB_PATH="/usr/local/cuda/bin/cuda-gdb"
+else
+    # Try glob for versioned CUDA installs
+    for p in /usr/local/cuda-*/bin/cuda-gdb; do
+        if [ -x "$p" ]; then
+            CUDA_GDB_PATH="$p"
+        fi
+    done
+fi
+if [ -z "$CUDA_GDB_PATH" ]; then
+    echo "  WARNING: cuda-gdb not found — using placeholder in launch.json"
+    CUDA_GDB_PATH="cuda-gdb"
+else
+    echo "  cuda-gdb: $CUDA_GDB_PATH"
+fi
+
+# Detect GPU SM architecture
+GPU_ARCH=""
+if command -v nvidia-smi &>/dev/null; then
+    COMPUTE_CAP=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null | head -1 | tr -d '[:space:]')
+    if [ -n "$COMPUTE_CAP" ]; then
+        GPU_ARCH="sm_$(echo "$COMPUTE_CAP" | tr -d '.')"
+    fi
+fi
+if [ -z "$GPU_ARCH" ]; then
+    echo "  WARNING: Could not detect GPU arch — defaulting to sm_75"
+    GPU_ARCH="sm_75"
+else
+    echo "  GPU arch: $GPU_ARCH"
+fi
+
+# Copy example CUDA source into workspace
+cp /tmp/saxpy.cu ~/workspace/saxpy.cu
+echo "  Deployed saxpy.cu"
+
+# Deploy launch.json with cuda-gdb path
+sed "s|__CUDA_GDB_PATH__|${CUDA_GDB_PATH}|g" /tmp/launch.json > ~/workspace/.vscode/launch.json
+echo "  Deployed launch.json"
+
+# Deploy tasks.json with GPU architecture
+sed "s|__GPU_ARCH__|${GPU_ARCH}|g" /tmp/tasks.json > ~/workspace/.vscode/tasks.json
+echo "  Deployed tasks.json"
+
 echo ""
 echo "=== Remote setup complete ==="

aws_bootstrap/resources/saxpy.cu

@@ -0,0 +1,49 @@
+/**
+ * SAXPY Example, CUDA Style
+ * Source: https://developer.nvidia.com/blog/easy-introduction-cuda-c-and-c/
+ *
+ * This is included as an example CUDA C++ source file to try out the VS Code launch configuration we include on the host machine.
+ *
+ */
+#include <stdio.h>
+
+__global__
+void saxpy(int n, float a, float *x, float *y)
+{
+  int i = blockIdx.x*blockDim.x + threadIdx.x;
+  if (i < n) y[i] = a*x[i] + y[i];
+}
+
+int main(void)
+{
+  int N = 1<<20;
+  float *x, *y, *d_x, *d_y;
+  x = (float*)malloc(N*sizeof(float));
+  y = (float*)malloc(N*sizeof(float));
+
+  cudaMalloc(&d_x, N*sizeof(float));
+  cudaMalloc(&d_y, N*sizeof(float));
+
+  for (int i = 0; i < N; i++) {
+    x[i] = 1.0f;
+    y[i] = 2.0f;
+  }
+
+  cudaMemcpy(d_x, x, N*sizeof(float), cudaMemcpyHostToDevice);
+  cudaMemcpy(d_y, y, N*sizeof(float), cudaMemcpyHostToDevice);
+
+  // Perform SAXPY on 1M elements
+  saxpy<<<(N+255)/256, 256>>>(N, 2.0f, d_x, d_y);
+
+  cudaMemcpy(y, d_y, N*sizeof(float), cudaMemcpyDeviceToHost);
+
+  float maxError = 0.0f;
+  for (int i = 0; i < N; i++)
+    maxError = max(maxError, abs(y[i]-4.0f));
+  printf("Max error: %f\n", maxError);
+
+  cudaFree(d_x);
+  cudaFree(d_y);
+  free(x);
+  free(y);
+}

aws_bootstrap/resources/tasks.json

@@ -0,0 +1,48 @@
+{
+    // CUDA build tasks for VSCode
+    // Deployed to: ~/workspace/.vscode/tasks.json
+    "version": "2.0.0",
+    "tasks": [
+        {
+            "label": "nvcc: build active file (debug)",
+            "type": "shell",
+            "command": "nvcc",
+            "args": [
+                "-g",                           // Host debug symbols
+                "-G",                           // Device (GPU) debug symbols
+                "-O0",                          // No optimization
+                "-arch=__GPU_ARCH__",            // GPU arch (auto-detected)
+                "-o",
+                "${fileDirname}/${fileBasenameNoExtension}",
+                "${file}"
+            ],
+            "options": {
+                "cwd": "${fileDirname}"
+            },
+            "problemMatcher": ["$nvcc"],
+            "group": {
+                "kind": "build",
+                "isDefault": true
+            },
+            "detail": "Compile active .cu file with debug symbols (-g -G)"
+        },
+        {
+            "label": "nvcc: build active file (release)",
+            "type": "shell",
+            "command": "nvcc",
+            "args": [
+                "-O3",
+                "-arch=__GPU_ARCH__",
+                "-o",
+                "${fileDirname}/${fileBasenameNoExtension}",
+                "${file}"
+            ],
+            "options": {
+                "cwd": "${fileDirname}"
+            },
+            "problemMatcher": ["$nvcc"],
+            "group": "build",
+            "detail": "Compile active .cu file optimized (no debug)"
+        }
+    ]
+}