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

aws_bootstrap/resources/gpu_smoke_test.ipynb

@@ -0,0 +1,340 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "title",
+   "metadata": {},
+   "source": [
+    "# GPU Smoke Test\n",
+    "\n",
+    "Interactive GPU verification and quick benchmark for aws-bootstrap instances.\n",
+    "\n",
+    "Run each cell top-to-bottom to verify the CUDA stack, exercise FP32/FP16 operations,\n",
+    "train a small CNN on MNIST, and visualise loss and memory usage.\n",
+    "\n",
+    "For the full CLI benchmark (CNN + Transformer, configurable precision/batch-size),\n",
+    "run `python ~/gpu_benchmark.py` from a terminal."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "system-info",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "\n",
+    "import torch\n",
+    "\n",
+    "\n",
+    "print(f\"Python  : {sys.version.split()[0]}\")\n",
+    "print(f\"PyTorch : {torch.__version__}\")\n",
+    "print(f\"CUDA    : {torch.version.cuda}\")\n",
+    "print(f\"cuDNN   : {torch.backends.cudnn.version()}\")\n",
+    "\n",
+    "assert torch.cuda.is_available(), \"CUDA is not available!\"\n",
+    "\n",
+    "props = torch.cuda.get_device_properties(0)\n",
+    "print(f\"\\nGPU           : {props.name}\")\n",
+    "print(f\"Compute cap.  : {props.major}.{props.minor}\")\n",
+    "print(f\"Total memory  : {props.total_memory / (1024**3):.1f} GB\")\n",
+    "print(f\"SM count      : {props.multi_processor_count}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "smoke-header",
+   "metadata": {},
+   "source": [
+    "## CUDA Smoke Tests"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "matmul-test",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "\n",
+    "\n",
+    "N = 1024\n",
+    "\n",
+    "# --- FP32 matmul ---\n",
+    "a32 = torch.randn(N, N, device=\"cuda\")\n",
+    "b32 = torch.randn(N, N, device=\"cuda\")\n",
+    "\n",
+    "start = torch.cuda.Event(enable_timing=True)\n",
+    "end = torch.cuda.Event(enable_timing=True)\n",
+    "\n",
+    "start.record()\n",
+    "c32 = torch.mm(a32, b32)\n",
+    "end.record()\n",
+    "torch.cuda.synchronize()\n",
+    "fp32_ms = start.elapsed_time(end)\n",
+    "print(f\"FP32 matmul ({N}x{N}): {fp32_ms:.2f} ms\")\n",
+    "\n",
+    "# --- FP16 matmul ---\n",
+    "a16 = a32.half()\n",
+    "b16 = b32.half()\n",
+    "\n",
+    "start.record()\n",
+    "c16 = torch.mm(a16, b16)\n",
+    "end.record()\n",
+    "torch.cuda.synchronize()\n",
+    "fp16_ms = start.elapsed_time(end)\n",
+    "print(f\"FP16 matmul ({N}x{N}): {fp16_ms:.2f} ms\")\n",
+    "\n",
+    "# Correctness check: FP16 result should be close to FP32\n",
+    "diff = (c32 - c16.float()).abs().max().item()\n",
+    "print(f\"Max abs diff FP32 vs FP16: {diff:.4f}\")\n",
+    "assert diff < N, f\"Unexpectedly large diff: {diff}\"  # loose bound\n",
+    "print(\"PASSED\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "amp-test",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "\n",
+    "\n",
+    "# AMP autocast: Linear + Conv2d\n",
+    "linear = nn.Linear(512, 512).cuda()\n",
+    "conv = nn.Conv2d(3, 64, 3, padding=1).cuda()\n",
+    "\n",
+    "x_lin = torch.randn(32, 512, device=\"cuda\")\n",
+    "x_conv = torch.randn(4, 3, 32, 32, device=\"cuda\")\n",
+    "\n",
+    "with torch.amp.autocast(device_type=\"cuda\"):\n",
+    "    y_lin = linear(x_lin)\n",
+    "    y_conv = conv(x_conv)\n",
+    "\n",
+    "torch.cuda.synchronize()\n",
+    "print(f\"Linear  output: {y_lin.shape}, dtype={y_lin.dtype}\")\n",
+    "print(f\"Conv2d  output: {y_conv.shape}, dtype={y_conv.dtype}\")\n",
+    "print(\"AMP autocast PASSED\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "memory-test",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "\n",
+    "\n",
+    "torch.cuda.empty_cache()\n",
+    "torch.cuda.reset_peak_memory_stats()\n",
+    "\n",
+    "before = torch.cuda.memory_allocated()\n",
+    "big = torch.randn(4096, 4096, device=\"cuda\")  # ~64 MB\n",
+    "allocated = torch.cuda.memory_allocated()\n",
+    "total = torch.cuda.get_device_properties(0).total_memory\n",
+    "\n",
+    "print(f\"Before alloc : {before / 1e6:.1f} MB\")\n",
+    "print(f\"After alloc  : {allocated / 1e6:.1f} MB\")\n",
+    "print(f\"Total GPU mem: {total / 1e9:.1f} GB\")\n",
+    "\n",
+    "del big\n",
+    "torch.cuda.empty_cache()\n",
+    "after_free = torch.cuda.memory_allocated()\n",
+    "print(f\"After free   : {after_free / 1e6:.1f} MB\")\n",
+    "assert after_free <= before + 1e6, \"Memory not freed!\"\n",
+    "print(\"Memory alloc/free PASSED\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "train-header",
+   "metadata": {},
+   "source": [
+    "## Quick Training Benchmark\n",
+    "\n",
+    "Train a small CNN on MNIST for 5 epochs and collect the loss per batch."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "mnist-train",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "from torch.utils.data import DataLoader\n",
+    "from torchvision import datasets, transforms\n",
+    "\n",
+    "\n",
+    "class MNISTConvNet(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.features = nn.Sequential(\n",
+    "            nn.Conv2d(1, 64, 3, padding=1),\n",
+    "            nn.BatchNorm2d(64),\n",
+    "            nn.ReLU(inplace=True),\n",
+    "            nn.Conv2d(64, 64, 3, padding=1),\n",
+    "            nn.BatchNorm2d(64),\n",
+    "            nn.ReLU(inplace=True),\n",
+    "            nn.MaxPool2d(2),\n",
+    "            nn.Conv2d(64, 128, 3, padding=1),\n",
+    "            nn.BatchNorm2d(128),\n",
+    "            nn.ReLU(inplace=True),\n",
+    "            nn.Conv2d(128, 128, 3, padding=1),\n",
+    "            nn.BatchNorm2d(128),\n",
+    "            nn.ReLU(inplace=True),\n",
+    "            nn.MaxPool2d(2),\n",
+    "            nn.Conv2d(128, 256, 3, padding=1),\n",
+    "            nn.BatchNorm2d(256),\n",
+    "            nn.ReLU(inplace=True),\n",
+    "            nn.AdaptiveAvgPool2d(1),\n",
+    "        )\n",
+    "        self.classifier = nn.Sequential(\n",
+    "            nn.Flatten(),\n",
+    "            nn.Linear(256, 256),\n",
+    "            nn.ReLU(inplace=True),\n",
+    "            nn.Dropout(0.5),\n",
+    "            nn.Linear(256, 10),\n",
+    "        )\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        return self.classifier(self.features(x))\n",
+    "\n",
+    "\n",
+    "device = torch.device(\"cuda\")\n",
+    "torch.cuda.reset_peak_memory_stats()\n",
+    "\n",
+    "transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])\n",
+    "train_set = datasets.MNIST(\"/tmp/data\", train=True, download=True, transform=transform)\n",
+    "loader = DataLoader(train_set, batch_size=256, shuffle=True, num_workers=2, pin_memory=True)\n",
+    "\n",
+    "model = MNISTConvNet().to(device)\n",
+    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n",
+    "scaler = torch.amp.GradScaler(\"cuda\")\n",
+    "\n",
+    "NUM_EPOCHS = 5\n",
+    "losses = []\n",
+    "\n",
+    "for epoch in range(1, NUM_EPOCHS + 1):\n",
+    "    model.train()\n",
+    "    epoch_loss = 0.0\n",
+    "    for images, labels in loader:\n",
+    "        images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True)\n",
+    "        optimizer.zero_grad(set_to_none=True)\n",
+    "        with torch.amp.autocast(device_type=\"cuda\"):\n",
+    "            loss = F.cross_entropy(model(images), labels)\n",
+    "        scaler.scale(loss).backward()\n",
+    "        scaler.step(optimizer)\n",
+    "        scaler.update()\n",
+    "        losses.append(loss.item())\n",
+    "        epoch_loss += loss.item()\n",
+    "    avg = epoch_loss / len(loader)\n",
+    "    print(f\"Epoch {epoch}/{NUM_EPOCHS}  avg loss: {avg:.4f}\")\n",
+    "\n",
+    "peak_mb = torch.cuda.max_memory_allocated() / (1024**2)\n",
+    "print(f\"\\nPeak GPU memory during training: {peak_mb:.0f} MB\")\n",
+    "print(f\"Total batches: {len(losses)}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "loss-plot",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "\n",
+    "plt.figure(figsize=(10, 4))\n",
+    "plt.plot(losses, linewidth=0.8, alpha=0.7)\n",
+    "plt.xlabel(\"Batch\")\n",
+    "plt.ylabel(\"Cross-Entropy Loss\")\n",
+    "plt.title(\"MNIST CNN Training Loss\")\n",
+    "plt.grid(True, alpha=0.3)\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "mem-header",
+   "metadata": {},
+   "source": [
+    "## GPU Memory & Utilization"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "mem-chart",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import torch\n",
+    "\n",
+    "\n",
+    "peak_mb = torch.cuda.max_memory_allocated() / (1024**2)\n",
+    "total_mb = torch.cuda.get_device_properties(0).total_memory / (1024**2)\n",
+    "free_mb = total_mb - peak_mb\n",
+    "\n",
+    "fig, ax = plt.subplots(figsize=(6, 4))\n",
+    "bars = ax.bar([\"Peak Used\", \"Remaining\"], [peak_mb, free_mb], color=[\"#e74c3c\", \"#2ecc71\"])\n",
+    "ax.set_ylabel(\"MB\")\n",
+    "ax.set_title(f\"GPU Memory: {peak_mb:.0f} MB peak / {total_mb:.0f} MB total\")\n",
+    "for bar in bars:\n",
+    "    ax.text(\n",
+    "        bar.get_x() + bar.get_width() / 2,\n",
+    "        bar.get_height() + 50,\n",
+    "        f\"{bar.get_height():.0f}\",\n",
+    "        ha=\"center\",\n",
+    "        va=\"bottom\",\n",
+    "        fontsize=11,\n",
+    "    )\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "summary",
+   "metadata": {},
+   "source": [
+    "## Summary\n",
+    "\n",
+    "If all cells above ran without error, the CUDA stack is healthy and the GPU is\n",
+    "ready for training workloads.\n",
+    "\n",
+    "### Next steps\n",
+    "\n",
+    "- **Full benchmark** (CNN + Transformer, configurable precision): `python ~/gpu_benchmark.py`\n",
+    "- **Jupyter tips**: use `!nvidia-smi` in a cell to check GPU utilisation at any time\n",
+    "- **VSCode Remote SSH**: connect with `ssh aws-gpu<N>` for a full IDE experience"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

aws_bootstrap/resources/remote_setup.sh

@@ -0,0 +1,188 @@
+#!/usr/bin/env bash
+# remote_setup.sh — Post-boot setup for Deep Learning AMI instances.
+# Runs on the EC2 instance after SSH becomes available.
+set -euo pipefail
+
+echo "=== aws-bootstrap-g4dn remote setup ==="
+
+# 1. Verify GPU
+echo ""
+echo "[1/5] 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
+    echo "WARNING: nvidia-smi not found"
+fi
+
+if command -v nvcc &>/dev/null; then
+    nvcc --version | grep "release"
+else
+    echo "WARNING: nvcc not found (CUDA toolkit may not be installed)"
+fi
+
+# 2. Install utilities
+echo ""
+echo "[2/5] 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..."
+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
+
+# --- CUDA-aware PyTorch installation ---
+# Known PyTorch CUDA wheel tags (ascending order).
+# Update this list when PyTorch publishes new CUDA builds.
+# See: https://download.pytorch.org/whl/
+KNOWN_CUDA_TAGS=(118 121 124 126 128 129 130)
+
+detect_cuda_version() {
+    # Primary: nvcc (actual toolkit installed on the system)
+    if command -v nvcc &>/dev/null; then
+        nvcc --version | grep -oP 'release \K[\d.]+'
+        return
+    fi
+    # Fallback: nvidia-smi (max CUDA the driver supports)
+    if command -v nvidia-smi &>/dev/null; then
+        nvidia-smi | grep -oP 'CUDA Version: \K[\d.]+'
+        return
+    fi
+    echo ""
+}
+
+cuda_version_to_tag() {
+    # "12.9" → "129", "13.0" → "130"
+    echo "$1" | tr -d '.'
+}
+
+find_best_cuda_tag() {
+    local detected_tag="$1"
+    local best=""
+    for tag in "${KNOWN_CUDA_TAGS[@]}"; do
+        if [ "$tag" -le "$detected_tag" ]; then
+            best="$tag"
+        fi
+    done
+    echo "$best"
+}
+
+install_pytorch_cuda() {
+    local cuda_ver
+    cuda_ver=$(detect_cuda_version)
+
+    if [ -z "$cuda_ver" ]; then
+        echo "  WARNING: No CUDA detected — installing PyTorch from PyPI (CPU or default CUDA)"
+        uv pip install --python ~/venv/bin/python torch torchvision
+        return
+    fi
+    echo "  Detected CUDA version: $cuda_ver"
+
+    local detected_tag
+    detected_tag=$(cuda_version_to_tag "$cuda_ver")
+
+    local best_tag
+    best_tag=$(find_best_cuda_tag "$detected_tag")
+
+    if [ -z "$best_tag" ]; then
+        echo "  WARNING: No matching PyTorch CUDA tag for cu${detected_tag} — installing from PyPI"
+        uv pip install --python ~/venv/bin/python torch torchvision
+        return
+    fi
+
+    echo "  Using PyTorch CUDA index: cu${best_tag}"
+    if ! uv pip install --python ~/venv/bin/python \
+            --default-index "https://download.pytorch.org/whl/cu${best_tag}" \
+            torch torchvision; then
+        echo "  WARNING: CUDA index install failed — falling back to PyPI"
+        uv pip install --python ~/venv/bin/python torch torchvision
+    fi
+}
+
+install_pytorch_cuda
+
+# Install remaining dependencies (torch/torchvision already installed above)
+uv pip install --python ~/venv/bin/python -r /tmp/requirements.txt
+
+# Copy GPU benchmark script and smoke test notebook
+cp /tmp/gpu_benchmark.py ~/gpu_benchmark.py
+cp /tmp/gpu_smoke_test.ipynb ~/gpu_smoke_test.ipynb
+
+# Auto-activate venv on login
+if ! grep -q 'source ~/venv/bin/activate' ~/.bashrc 2>/dev/null; then
+    echo 'source ~/venv/bin/activate' >> ~/.bashrc
+fi
+
+# Quick CUDA smoke test
+echo "  Running CUDA smoke test..."
+if ~/venv/bin/python -c "
+import torch
+assert torch.cuda.is_available(), 'CUDA not available'
+x = torch.randn(256, 256, device='cuda')
+y = torch.mm(x, x)
+torch.cuda.synchronize()
+print(f'  PyTorch {torch.__version__}, CUDA {torch.version.cuda}, GPU: {torch.cuda.get_device_name(0)}')
+print('  Quick matmul test: PASSED')
+"; then
+    echo "  CUDA smoke test passed"
+else
+    echo "  WARNING: CUDA smoke test failed — check PyTorch/CUDA installation"
+fi
+
+JUPYTER_CONFIG_DIR="$HOME/.jupyter"
+mkdir -p "$JUPYTER_CONFIG_DIR"
+cat > "$JUPYTER_CONFIG_DIR/jupyter_lab_config.py" << 'PYEOF'
+c.ServerApp.ip = '0.0.0.0'
+c.ServerApp.port = 8888
+c.ServerApp.open_browser = False
+c.IdentityProvider.token = ''
+c.ServerApp.allow_remote_access = True
+PYEOF
+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..."
+LOGIN_USER=$(whoami)
+
+sudo tee /etc/systemd/system/jupyter.service > /dev/null << SVCEOF
+[Unit]
+Description=Jupyter Lab Server
+After=network.target
+
+[Service]
+Type=simple
+User=${LOGIN_USER}
+WorkingDirectory=/home/${LOGIN_USER}
+ExecStart=/home/${LOGIN_USER}/venv/bin/python -m jupyterlab
+Restart=on-failure
+RestartSec=10
+
+[Install]
+WantedBy=multi-user.target
+SVCEOF
+
+sudo systemctl daemon-reload
+sudo systemctl enable jupyter.service
+sudo systemctl start jupyter.service
+echo "  Jupyter service started (port 8888)"
+
+# 5. SSH keepalive
+echo ""
+echo "[5/5] 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
+    sudo systemctl reload sshd
+    echo "  SSH keepalive configured"
+else
+    echo "  SSH keepalive already configured"
+fi
+
+echo ""
+echo "=== Remote setup complete ==="

aws_bootstrap/resources/requirements.txt

@@ -0,0 +1,8 @@
+numpy
+matplotlib
+scipy
+datasets
+transformers
+pandas
+tensorboard
+jupyterlab